Repository: risingwavelabs/risingwave (language: python)

Issue:
refactor(frontend): rework `UPDATE` to support subqueries
I hereby agree to the terms of the [RisingWave Labs, Inc. Contributor License Agreement](https://gist.github.com/TennyZhuang/f00be7f16996ea48effb049aa7be4d66#file-rw_cla).

## What's changed and what's your intention?

Rework `UPDATE` to support subqueries. Close #2672.

----

Previously, `Update` has a list of `index -> expr` to calculate the updated value for each column. This means that its capability cannot exceed a single `Project` node, so that there's no way to support subquery assignments like `UPDATE (v1, v2) = (select 114, 514);`.

After this PR, the input of an `Update` will calculate the old and new values for all columns, so that `Update` only needs to pick some columns with `InputRef` to generate the `U-` and `U+` chunks. Complex assignments like subqueries will be handled by the optimizer, just like within a normal `SELECT` statement.

To support this, a new subquery kind named `UpdateSet` is added to allow returning multiple columns (still single row) by compositing them into a `struct` type.

An example demonstrating the plan:

```sql
dev=> explain update t set (v1, v2) = (select v, 666 from generate_series(1, 1) as g(v));
                                          QUERY PLAN                                          
----------------------------------------------------------------------------------------------
 BatchExchange { order: [], dist: Single }
 └─BatchUpdate { table: t, exprs: [Field($4, 0:Int32), Field($4, 1:Int32), $2] }
   └─BatchNestedLoopJoin { type: LeftOuter, predicate: true }
     ├─BatchExchange { order: [], dist: Single }
     │ └─BatchScan { table: t, columns: [v1, v2, _row_id, _rw_timestamp] }
     └─BatchMaxOneRow
       └─BatchProject { exprs: [Row(GenerateSeries(1:Int32, 1:Int32), 666:Int32) as $expr1] }
         └─BatchProjectSet { select_list: [GenerateSeries(1:Int32, 1:Int32)] }
           └─BatchValues { rows: [[]] }
(9 rows)
```

## Checklist

- [x] I have written necessary rustdoc comments
- [x] I have added necessary unit tests and integration tests
- [x] All checks passed in `./risedev check` (or alias, `./risedev c`)

## Documentation

- [ ] My PR needs documentation updates. (Please use the **Release note** section below to summarize the impact on users)

## Release note

If this PR includes changes that directly affect users or other significant modifications relevant to the community, kindly draft a release note to provide a concise summary of these changes. Please prioritize highlighting the impact these changes will have on users.


<!--
Please create a release note for your changes.

Discuss technical details in the "What's changed" section, and
focus on the impact on users in the release note.

You should also mention the environment or conditions where the impact may occur.
-->

</details>

refactor(optimizer): record error contexts when casting composite types
I hereby agree to the terms of the [RisingWave Labs, Inc. Contributor License Agreement](https://gist.github.com/TennyZhuang/f00be7f16996ea48effb049aa7be4d66#file-rw_cla).

## What's changed and what's your intention?

Record contexts for error when recursively calling `cast` on composite types, to provide better error messages.

Preview:

https://github.com/risingwavelabs/risingwave/blob/fc4223348295f029a4d88d0d9f64b4de8a760ac8/src/frontend/planner_test/tests/testdata/output/cast.yaml#L92-L115

## Checklist

- [x] I have written necessary rustdoc comments
- [x] I have added necessary unit tests and integration tests
- [x] All checks passed in `./risedev check` (or alias, `./risedev c`)

## Documentation

- [ ] My PR needs documentation updates. (Please use the **Release note** section below to summarize the impact on users)

## Release note

If this PR includes changes that directly affect users or other significant modifications relevant to the community, kindly draft a release note to provide a concise summary of these changes. Please prioritize highlighting the impact these changes will have on users.


<!--
Please create a release note for your changes.

Discuss technical details in the "What's changed" section, and
focus on the impact on users in the release note.

You should also mention the environment or conditions where the impact may occur.
-->

</details>



Discussion / hints:
> A corner case

In this case, the query `select 1, (select y.a from y)` should also be rejected, but we accepted it. So I presume there's something wrong with `BatchMaxOneRow`. 🤔

I think it's because we don't enforce input of `BatchMaxOneRow` to be singleton. 🥵
> I think it's because we don't enforce input of `BatchMaxOneRow` to be singleton

Oh, that cloud be the problem
> > I think it's because we don't enforce input of `BatchMaxOneRow` to be singleton
> 
> Oh, that cloud be the problem

Will fix in https://github.com/risingwavelabs/risingwave/pull/19452


Fix patch:
diff --git a/src/frontend/src/binder/update.rs b/src/frontend/src/binder/update.rs
index f57ad1d197982..a2038a4d471d9 100644
--- a/src/frontend/src/binder/update.rs
+++ b/src/frontend/src/binder/update.rs
@@ -161,8 +161,8 @@ impl Binder {
 
         for Assignment { id, value } in assignments {
             let ids: Vec<_> = id
-                .into_iter()
-                .map(|id| self.bind_expr(Expr::Identifier(id)))
+                .iter()
+                .map(|id| self.bind_expr(Expr::Identifier(id.clone())))
                 .try_collect()?;
 
             match (ids.as_slice(), value) {
@@ -206,8 +206,9 @@ impl Binder {
                         bail_bind_error!("number of columns does not match number of values");
                     }
 
-                    let target_type = DataType::new_unnamed_struct(
+                    let target_type = DataType::new_struct(
                         ids.iter().map(|id| id.return_type()).collect(),
+                        id.iter().map(|id| id.real_value()).collect(),
                     );
                     let expr = expr.cast_assign(target_type)?;
 
diff --git a/src/frontend/src/expr/subquery.rs b/src/frontend/src/expr/subquery.rs
index 8460f73d5fbba..fc904638790e2 100644
--- a/src/frontend/src/expr/subquery.rs
+++ b/src/frontend/src/expr/subquery.rs
@@ -14,10 +14,10 @@
 
 use std::hash::Hash;
 
-use risingwave_common::types::DataType;
+use risingwave_common::types::{DataType, StructType};
 
 use super::{Expr, ExprImpl, ExprType};
-use crate::binder::BoundQuery;
+use crate::binder::{BoundQuery, UNNAMED_COLUMN};
 use crate::expr::{CorrelatedId, Depth};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -91,7 +91,19 @@ impl Expr for Subquery {
                 assert_eq!(types.len(), 1, "Subquery with more than one column");
                 types[0].clone()
             }
-            SubqueryKind::UpdateSet => DataType::new_unnamed_struct(self.query.data_types()),
+            SubqueryKind::UpdateSet => {
+                let schema = self.query.schema();
+                let struct_type = if schema.fields().iter().any(|f| f.name == UNNAMED_COLUMN) {
+                    StructType::unnamed(self.query.data_types())
+                } else {
+                    StructType::new(
+                        (schema.fields().iter().cloned())
+                            .map(|f| (f.name, f.data_type))
+                            .collect(),
+                    )
+                };
+                DataType::Struct(struct_type)
+            }
             SubqueryKind::Array => {
                 let types = self.query.data_types();
                 assert_eq!(types.len(), 1, "Subquery with more than one column");
diff --git a/src/frontend/src/expr/type_inference/cast.rs b/src/frontend/src/expr/type_inference/cast.rs
index c9b09fe18ecac..e03ec98883ce9 100644
--- a/src/frontend/src/expr/type_inference/cast.rs
+++ b/src/frontend/src/expr/type_inference/cast.rs
@@ -206,22 +206,13 @@ fn cast_struct(source: &DataType, target: &DataType, allows: CastContext) -> Cas
                         Ok(())
                     } else {
                         cast(src_ty, dst_ty, allows).map_err(|inner| {
-                            if src_name.is_empty() {
-                                inner
-                            } else if dst_name.is_empty() {
-                                cast_error!(
-                                    source = inner,
-                                    "cannot cast struct field \"{}\"",
-                                    src_name
-                                )
-                            } else {
-                                cast_error!(
-                                    source = inner,
-                                    "cannot cast struct field \"{}\" to struct field \"{}\"",
-                                    src_name,
-                                    dst_name
-                                )
-                            }
+                            let cast_from = (!src_name.is_empty())
+                                .then(|| format!(" struct field \"{}\"", src_name))
+                                .unwrap_or_default();
+                            let cast_to = (!dst_name.is_empty())
+                                .then(|| format!(" to struct field \"{}\"", dst_name))
+                                .unwrap_or_default();
+                            cast_error!(source = inner, "cannot cast{}{}", cast_from, cast_to)
                         })
                     }
                 },


Test patch:
diff --git a/e2e_test/batch/basic/dml_update.slt.part b/e2e_test/batch/basic/dml_update.slt.part
index fc2647cea147b..0adc1ce5f27e6 100644
--- a/e2e_test/batch/basic/dml_update.slt.part
+++ b/e2e_test/batch/basic/dml_update.slt.part
@@ -99,8 +99,21 @@ update t set (v1, v2) = (select '888.88', 999);
 db error: ERROR: Failed to run the query
 
 Caused by these errors (recent errors listed first):
-  1: cannot cast type "record" to "record"
-  2: cannot cast type "character varying" to "integer" in Assign context
+  1: cannot cast type "record" to "struct<v1 integer, v2 integer>"
+  2: cannot cast to struct field "v1"
+  3: cannot cast type "character varying" to "integer" in Assign context
+
+
+# Multiple assignments, to subquery (with column name) with cast failure.
+statement error
+update t set (v1, v2) = (select '888.88' s1, 999 s2);
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: cannot cast type "struct<s1 character varying, s2 integer>" to "struct<v1 integer, v2 integer>"
+  2: cannot cast struct field "s1" to struct field "v1"
+  3: cannot cast type "character varying" to "integer" in Assign context
 
 
 # Multiple assignments, to subquery with mismatched column count.
diff --git a/src/frontend/planner_test/tests/testdata/output/update.yaml b/src/frontend/planner_test/tests/testdata/output/update.yaml
index 26c6d52dc5e05..42246cabbb0f1 100644
--- a/src/frontend/planner_test/tests/testdata/output/update.yaml
+++ b/src/frontend/planner_test/tests/testdata/output/update.yaml
@@ -235,7 +235,7 @@
   batch_plan: |-
     BatchExchange { order: [], dist: Single }
     └─BatchUpdate { table: t, exprs: [Field($4, 0:Int32), Field($4, 1:Int32), $2] }
-      └─BatchProject { exprs: [t.v1, t.v2, t._row_id, t._rw_timestamp, $expr10011::Struct(StructType { field_names: [], field_types: [Int32, Int32] }) as $expr1] }
+      └─BatchProject { exprs: [t.v1, t.v2, t._row_id, t._rw_timestamp, $expr10011::Struct(StructType { field_names: ["v1", "v2"], field_types: [Int32, Int32] }) as $expr1] }
         └─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
           ├─BatchExchange { order: [], dist: Single }
           │ └─BatchScan { table: t, columns: [t.v1, t.v2, t._row_id, t._rw_timestamp], distribution: UpstreamHashShard(t._row_id) }