Repository: risingwavelabs/risingwave (language: python)

Issue:
feat(optimizer): runtime check for scalar subquery in batch queries
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?

In #12908, a hack is introduced to allow a system query to be executed. We concluded that the support for unique keys may help to get rid of it. However, as discussed in #5019 and #1335, it appears that supporting runtime check for scalar subquery can be much easier to implement.

~~I reuse the plan nodes and executor for `Limit` to achieve that, not sure if it's a good practice.~~

Introduce a new plan node named `MaxOneRow` and its corresponding batch executor for doing this.

## 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:
## [Codecov](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) Report
Attention: `78 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`1cacc07`)](https://app.codecov.io/gh/risingwavelabs/risingwave/commit/1cacc07fc6c224f0afe3f8da7ffaa4326d38bfa9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) 68.07% compared to head [(`46e2f6b`)](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) 68.00%.
> Report is 13 commits behind head on main.

| [Files](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) | Patch % | Lines |
|---|---|---|
| [...ntend/src/optimizer/plan\_node/batch\_max\_one\_row.rs](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs#diff-c3JjL2Zyb250ZW5kL3NyYy9vcHRpbWl6ZXIvcGxhbl9ub2RlL2JhdGNoX21heF9vbmVfcm93LnJz) | 0.00% | [29 Missing :warning: ](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) |
| [...end/src/optimizer/plan\_node/logical\_max\_one\_row.rs](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs#diff-c3JjL2Zyb250ZW5kL3NyYy9vcHRpbWl6ZXIvcGxhbl9ub2RlL2xvZ2ljYWxfbWF4X29uZV9yb3cucnM=) | 55.55% | [24 Missing :warning: ](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) |
| [src/batch/src/executor/max\_one\_row.rs](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs#diff-c3JjL2JhdGNoL3NyYy9leGVjdXRvci9tYXhfb25lX3Jvdy5ycw==) | 65.00% | [21 Missing :warning: ](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) |
| [...end/src/optimizer/plan\_node/generic/max\_one\_row.rs](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs#diff-c3JjL2Zyb250ZW5kL3NyYy9vcHRpbWl6ZXIvcGxhbl9ub2RlL2dlbmVyaWMvbWF4X29uZV9yb3cucnM=) | 78.94% | [4 Missing :warning: ](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main   #13880      +/-   ##
==========================================
- Coverage   68.07%   68.00%   -0.07%     
==========================================
  Files        1548     1552       +4     
  Lines      267474   267653     +179     
==========================================
- Hits       182075   182018      -57     
- Misses      85399    85635     +236     
```

| [Flag](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) | Coverage Δ | |
|---|---|---|
| [rust](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs) | `68.00% <65.93%> (-0.07%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs#carryforward-flags-in-the-pull-request-comment) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/risingwavelabs/risingwave/pull/13880?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=risingwavelabs).

> What's your motivation for doing this?

It's #12908. 😄
I am considering whether we should use a dedicated operator e.g. `MaxOneRow` to abstract this semantic instead of coupling with the `Limit`. 
Pros: we can use `Limit` more easily without caring about the `exceeding_limit` restrictions. 
> I am considering whether we should use a dedicated operator e.g. `MaxOneRow` to abstract this semantic instead of coupling with the `Limit`. Pros: we can use `Limit` more easily without caring about the `exceeding_limit` restrictions.

Refactored the implementation. PTAL again. 🥰
> So the subquery is executed during the optimization phase?

Not exactly. We now enable the checking of scalar subquery cardinality at runtime. Previously, only compile-time information was used, which resulted in some queries being unable to be planned.

Fix patch:
diff --git a/src/frontend/src/optimizer/plan_node/batch_max_one_row.rs b/src/frontend/src/optimizer/plan_node/batch_max_one_row.rs
index e35d6976672c4..4b3d0de43cb66 100644
--- a/src/frontend/src/optimizer/plan_node/batch_max_one_row.rs
+++ b/src/frontend/src/optimizer/plan_node/batch_max_one_row.rs
@@ -25,6 +25,7 @@ use super::{
 use crate::error::Result;
 use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
 use crate::optimizer::plan_node::ToLocalBatch;
+use crate::optimizer::property::{Order, RequiredDist};
 
 /// [`BatchMaxOneRow`] fetches up to one row from the input, returning an error
 /// if the input contains more than one row at runtime.
@@ -66,7 +67,9 @@ impl Distill for BatchMaxOneRow {
 
 impl ToDistributedBatch for BatchMaxOneRow {
     fn to_distributed(&self) -> Result<PlanRef> {
-        Ok(self.clone_with_input(self.input().to_distributed()?).into())
+        let new_input = RequiredDist::single()
+            .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
+        Ok(self.clone_with_input(new_input).into())
     }
 }
 
@@ -78,7 +81,9 @@ impl ToBatchPb for BatchMaxOneRow {
 
 impl ToLocalBatch for BatchMaxOneRow {
     fn to_local(&self) -> Result<PlanRef> {
-        Ok(self.clone_with_input(self.input().to_local()?).into())
+        let new_input = RequiredDist::single()
+            .enforce_if_not_satisfies(self.input().to_local()?, &Order::any())?;
+        Ok(self.clone_with_input(new_input).into())
     }
 }
 


Test patch:
diff --git a/e2e_test/batch/distribution_mode.slt b/e2e_test/batch/distribution_mode.slt
index b680796277c13..f101e5c5446e3 100644
--- a/e2e_test/batch/distribution_mode.slt
+++ b/e2e_test/batch/distribution_mode.slt
@@ -13,6 +13,7 @@ include ./aggregate/*.slt.part
 include ./types/*.slt.part
 include ./functions/*.slt.part
 include ./over_window/main.slt.part
+include ./subquery/**/*.slt.part
 
 statement ok
 SET QUERY_MODE TO auto;
diff --git a/e2e_test/batch/local_mode.slt b/e2e_test/batch/local_mode.slt
index 68df9f0d91950..a64fd49b85c8f 100644
--- a/e2e_test/batch/local_mode.slt
+++ b/e2e_test/batch/local_mode.slt
@@ -14,6 +14,7 @@ include ./types/*.slt.part
 include ./catalog/*.slt.part
 include ./functions/*.slt.part
 include ./over_window/main.slt.part
+include ./subquery/**/*.slt.part
 
 statement ok
 SET QUERY_MODE TO auto;
diff --git a/e2e_test/batch/subquery/scalar_subquery.slt.part b/e2e_test/batch/subquery/scalar_subquery.slt.part
index a0676e98f245c..5bb08201b7a75 100644
--- a/e2e_test/batch/subquery/scalar_subquery.slt.part
+++ b/e2e_test/batch/subquery/scalar_subquery.slt.part
@@ -21,8 +21,10 @@ select (select x from t) x, 1 one;
 statement error Scalar subquery might produce more than one row
 create materialized view mv as select (select x from t) x, 1 one;
 
+# Use a random value here to occasionally make it not in the same shard as `114514`,
+# demonstrating that `BatchMaxOneRow` correctly handles distributed subqueries.
 statement ok
-insert into t values (1919810);
+insert into t values (extract(epoch from now()));
 
 # Cannot query as the cardinality of the subquery is now 2
 query error Scalar subquery produced more than one row
diff --git a/src/frontend/planner_test/tests/testdata/input/subquery_expr.yaml b/src/frontend/planner_test/tests/testdata/input/subquery_expr.yaml
index d5b6149f70877..3604ca4fa0d0b 100644
--- a/src/frontend/planner_test/tests/testdata/input/subquery_expr.yaml
+++ b/src/frontend/planner_test/tests/testdata/input/subquery_expr.yaml
@@ -9,6 +9,7 @@
   expected_outputs:
   - logical_plan
   - optimized_logical_plan_for_batch
+  - batch_plan
   - stream_error
 - sql: |
     create table t(x int);
@@ -28,6 +29,7 @@
   expected_outputs:
   - logical_plan
   - optimized_logical_plan_for_batch
+  - batch_plan
   - stream_error
 - sql: |
     create table t(x int);
@@ -35,6 +37,7 @@
   expected_outputs:
   - logical_plan
   - optimized_logical_plan_for_batch
+  - batch_plan
   - stream_error
 - sql: |
     create table t(x int);
@@ -42,6 +45,7 @@
   expected_outputs:
   - logical_plan
   - optimized_logical_plan_for_batch
+  - batch_plan
   - stream_error
 - sql: |
     create table t(x int);
@@ -49,6 +53,7 @@
   expected_outputs:
   - logical_plan
   - optimized_logical_plan_for_batch
+  - batch_plan
   - stream_error
 - sql: |
     create table t1 (x int, y int);
diff --git a/src/frontend/planner_test/tests/testdata/output/subquery_expr.yaml b/src/frontend/planner_test/tests/testdata/output/subquery_expr.yaml
index a43a1d2df7166..519b876c8c67d 100644
--- a/src/frontend/planner_test/tests/testdata/output/subquery_expr.yaml
+++ b/src/frontend/planner_test/tests/testdata/output/subquery_expr.yaml
@@ -26,6 +26,14 @@
       ├─LogicalScan { table: t, columns: [] }
       └─LogicalMaxOneRow
         └─LogicalScan { table: t, columns: [t.x] }
+  batch_plan: |-
+    BatchProject { exprs: [t.x, 1:Int32] }
+    └─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+      ├─BatchExchange { order: [], dist: Single }
+      │ └─BatchScan { table: t, columns: [], distribution: SomeShard }
+      └─BatchMaxOneRow
+        └─BatchExchange { order: [], dist: Single }
+          └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
   stream_error: Scalar subquery might produce more than one row.
 - sql: |
     create table t(x int);
@@ -76,6 +84,16 @@
       └─LogicalMaxOneRow
         └─LogicalTopN { order: [t.x ASC], limit: 1, offset: 0, with_ties: true }
           └─LogicalScan { table: t, columns: [t.x] }
+  batch_plan: |-
+    BatchProject { exprs: [t.x, 1:Int32] }
+    └─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+      ├─BatchExchange { order: [], dist: Single }
+      │ └─BatchScan { table: t, columns: [], distribution: SomeShard }
+      └─BatchMaxOneRow
+        └─BatchTopN { order: [t.x ASC], limit: 1, offset: 0, with_ties: true }
+          └─BatchExchange { order: [], dist: Single }
+            └─BatchTopN { order: [t.x ASC], limit: 1, offset: 0, with_ties: true }
+              └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
   stream_error: Scalar subquery might produce more than one row.
 - sql: |
     create table t(x int);
@@ -92,6 +110,14 @@
       ├─LogicalScan { table: t, columns: [] }
       └─LogicalMaxOneRow
         └─LogicalScan { table: t, columns: [t.x] }
+  batch_plan: |-
+    BatchProject { exprs: [(t.x + 1:Int32) as $expr1] }
+    └─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+      ├─BatchExchange { order: [], dist: Single }
+      │ └─BatchScan { table: t, columns: [], distribution: SomeShard }
+      └─BatchMaxOneRow
+        └─BatchExchange { order: [], dist: Single }
+          └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
   stream_error: Scalar subquery might produce more than one row.
 - sql: |
     create table t(x int);
@@ -112,6 +138,14 @@
     │ └─LogicalMaxOneRow
     │   └─LogicalScan { table: t, columns: [t.x] }
     └─LogicalValues { rows: [[1:Int32]], schema: Schema { fields: [1:Int32:Int32] } }
+  batch_plan: |-
+    BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+    ├─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+    │ ├─BatchValues { rows: [[]] }
+    │ └─BatchMaxOneRow
+    │   └─BatchExchange { order: [], dist: Single }
+    │     └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
+    └─BatchValues { rows: [[1:Int32]] }
   stream_error: Scalar subquery might produce more than one row.
 - sql: |
     create table t(x int);
@@ -135,6 +169,19 @@
             ├─LogicalScan { table: t, columns: [t.x] }
             └─LogicalMaxOneRow
               └─LogicalScan { table: t, columns: [t.x] }
+  batch_plan: |-
+    BatchProject { exprs: [(t.x + $expr1) as $expr2] }
+    └─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+      ├─BatchExchange { order: [], dist: Single }
+      │ └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
+      └─BatchMaxOneRow
+        └─BatchProject { exprs: [(t.x + t.x) as $expr1] }
+          └─BatchNestedLoopJoin { type: LeftOuter, predicate: true, output: all }
+            ├─BatchExchange { order: [], dist: Single }
+            │ └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
+            └─BatchMaxOneRow
+              └─BatchExchange { order: [], dist: Single }
+                └─BatchScan { table: t, columns: [t.x], distribution: SomeShard }
   stream_error: Scalar subquery might produce more than one row.
 - sql: |
     create table t1 (x int, y int);