feat(batch): distributed dml

src/common/src/session_config/mod.rs

@@ -94,6 +94,10 @@ pub struct ConfigMap {
     #[parameter(default = true, rename = "rw_batch_enable_sort_agg")]
     batch_enable_sort_agg: bool,
 
+    /// Enable distributed dml, so a insert, delete and update statement could be executed in distributed way (e.g. running in multiple compute nodes).
+    #[parameter(default = false, rename = "batch_enable_distributed_dml")]
+    batch_enable_distributed_dml: bool,
+
     /// The max gap allowed to transform small range scan into multi point lookup.
     #[parameter(default = 8)]
     max_split_range_gap: i32,

src/frontend/planner_test/src/lib.rs

@@ -67,6 +67,8 @@ pub enum TestType {
     BatchPlanProto,
     /// Batch plan for local execution `.gen_batch_local_plan()`
     BatchLocalPlan,
+    /// Batch plan for local execution `.gen_batch_distributed_plan()`
+    BatchDistributedPlan,
 
     /// Create MV plan `.gen_create_mv_plan()`
     StreamPlan,
@@ -201,6 +203,9 @@ pub struct TestCaseResult {
     /// Batch plan for local execution `.gen_batch_local_plan()`
     pub batch_local_plan: Option<String>,
 
+    /// Batch plan for distributed execution `.gen_batch_distributed_plan()`
+    pub batch_distributed_plan: Option<String>,
+
     /// Generate sink plan
     pub sink_plan: Option<String>,
 
@@ -709,6 +714,36 @@ impl TestCase {
             }
         }
 
+        'distributed_batch: {
+            if self
+                .expected_outputs
+                .contains(&TestType::BatchDistributedPlan)
+                || self.expected_outputs.contains(&TestType::BatchError)
+            {
+                let batch_plan = match logical_plan.gen_batch_plan() {
+                    Ok(batch_plan) => match logical_plan.gen_batch_distributed_plan(batch_plan) {
+                        Ok(batch_plan) => batch_plan,
+                        Err(err) => {
+                            ret.batch_error = Some(err.to_report_string_pretty());
+                            break 'distributed_batch;
+                        }
+                    },
+                    Err(err) => {
+                        ret.batch_error = Some(err.to_report_string_pretty());
+                        break 'distributed_batch;
+                    }
+                };
+
+                // Only generate batch_plan if it is specified in test case
+                if self
+                    .expected_outputs
+                    .contains(&TestType::BatchDistributedPlan)
+                {
+                    ret.batch_distributed_plan = Some(explain_plan(&batch_plan));
+                }
+            }
+        }
+
         {
             // stream
             for (

src/frontend/planner_test/tests/testdata/input/delete.yaml

@@ -18,3 +18,10 @@
     delete from t returning sum(a);
   expected_outputs:
   - binder_error
+- name: distributed delete
+  sql: |
+    set batch_enable_distributed_dml = true;
+    create table t (a int, b int);
+    delete from t;
+  expected_outputs:
+    - batch_distributed_plan

src/frontend/planner_test/tests/testdata/input/insert.yaml

@@ -257,3 +257,10 @@
   expected_outputs:
   - batch_plan
   - logical_plan
+- name: distributed insert
+  sql: |
+    set batch_enable_distributed_dml = true;
+    create table t (a int, b int);
+    insert into t select * from t;
+  expected_outputs:
+  - batch_distributed_plan

src/frontend/planner_test/tests/testdata/input/update.yaml

@@ -91,3 +91,10 @@
   expected_outputs:
     - logical_plan
     - batch_plan
+- name: distributed update
+  sql: |
+    set batch_enable_distributed_dml = true;
+    create table t (a int, b int);
+    update t set a = a + 1;
+  expected_outputs:
+    - batch_distributed_plan

src/frontend/planner_test/tests/testdata/output/delete.yaml

@@ -32,3 +32,14 @@
     create table t (a int, b int);
     delete from t returning sum(a);
   binder_error: 'Bind error: should not have agg/window in the `RETURNING` list'
+- name: distributed delete
+  sql: |
+    set batch_enable_distributed_dml = true;
+    create table t (a int, b int);
+    delete from t;
+  batch_distributed_plan: |-
+    BatchSimpleAgg { aggs: [sum()] }
+    └─BatchExchange { order: [], dist: Single }
+      └─BatchDelete { table: t }
+        └─BatchExchange { order: [], dist: HashShard(t.a, t.b, t._row_id) }
+          └─BatchScan { table: t, columns: [t.a, t.b, t._row_id], distribution: UpstreamHashShard(t._row_id) }

src/frontend/planner_test/tests/testdata/output/insert.yaml

@@ -332,3 +332,14 @@
     BatchExchange { order: [], dist: Single }
     └─BatchInsert { table: t, mapping: [0:0], default: [1<-(2:Int32 + 3:Int32)] }
       └─BatchValues { rows: [[1:Int32]] }
+- name: distributed insert
+  sql: |
+    set batch_enable_distributed_dml = true;
+    create table t (a int, b int);
+    insert into t select * from t;
+  batch_distributed_plan: |-
+    BatchSimpleAgg { aggs: [sum()] }
+    └─BatchExchange { order: [], dist: Single }
+      └─BatchInsert { table: t, mapping: [0:0, 1:1] }
+        └─BatchExchange { order: [], dist: HashShard(t.a, t.b) }
+          └─BatchScan { table: t, columns: [t.a, t.b], distribution: SomeShard }

src/frontend/planner_test/tests/testdata/output/update.yaml

@@ -154,3 +154,14 @@
           │ └─BatchScan { table: t, columns: [t.a, t.b, t._row_id], distribution: UpstreamHashShard(t._row_id) }
           └─BatchExchange { order: [], dist: HashShard(t.b) }
             └─BatchScan { table: t, columns: [t.b], distribution: SomeShard }
+- name: distributed update
+  sql: |
+    set batch_enable_distributed_dml = true;
+    create table t (a int, b int);
+    update t set a = a + 1;
+  batch_distributed_plan: |-
+    BatchSimpleAgg { aggs: [sum()] }
+    └─BatchExchange { order: [], dist: Single }
+      └─BatchUpdate { table: t, exprs: [($0 + 1:Int32), $1, $2] }
+        └─BatchExchange { order: [], dist: HashShard(t.a, t.b, t._row_id) }
+          └─BatchScan { table: t, columns: [t.a, t.b, t._row_id], distribution: UpstreamHashShard(t._row_id) }

src/frontend/src/optimizer/plan_node/batch_delete.rs

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 use risingwave_common::error::Result;
+use risingwave_common::types::DataType;
+use risingwave_expr::aggregate::AggKind;
 use risingwave_pb::batch_plan::plan_node::NodeBody;
 use risingwave_pb::batch_plan::DeleteNode;
 
@@ -21,10 +23,12 @@ use super::utils::impl_distill_by_unit;
 use super::{
     generic, ExprRewritable, PlanBase, PlanRef, PlanTreeNodeUnary, ToBatchPb, ToDistributedBatch,
 };
+use crate::expr::InputRef;
 use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
-use crate::optimizer::plan_node::generic::PhysicalPlanRef;
-use crate::optimizer::plan_node::ToLocalBatch;
+use crate::optimizer::plan_node::generic::{Agg, GenericPlanNode, PhysicalPlanRef};
+use crate::optimizer::plan_node::{BatchSimpleAgg, PlanAggCall, ToLocalBatch};
 use crate::optimizer::property::{Distribution, Order, RequiredDist};
+use crate::utils::{Condition, IndexSet};
 
 /// `BatchDelete` implements [`super::LogicalDelete`]
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -35,7 +39,6 @@ pub struct BatchDelete {
 
 impl BatchDelete {
     pub fn new(core: generic::Delete<PlanRef>) -> Self {
-        assert_eq!(core.input.distribution(), &Distribution::Single);
         let base =
             PlanBase::new_batch_with_core(&core, core.input.distribution().clone(), Order::any());
         Self { base, core }
@@ -59,9 +62,43 @@ impl_distill_by_unit!(BatchDelete, core, "BatchDelete");
 
 impl ToDistributedBatch for BatchDelete {
     fn to_distributed(&self) -> Result<PlanRef> {
-        let new_input = RequiredDist::single()
+        if self
+            .core
+            .ctx()
+            .session_ctx()
+            .config()
+            .batch_enable_distributed_dml()
+        {
+            // Add an hash shuffle between the delete and its input.
+            let new_input = RequiredDist::PhysicalDist(Distribution::HashShard(
+                (0..self.input().schema().len()).collect(),
+            ))
             .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
-        Ok(self.clone_with_input(new_input).into())
+            let new_delete: PlanRef = self.clone_with_input(new_input).into();
+            if self.core.returning {
+                Ok(new_delete)
+            } else {
+                let new_delete =
+                    RequiredDist::single().enforce_if_not_satisfies(new_delete, &Order::any())?;
+                // Accumulate the affected rows.
+                let sum_agg = PlanAggCall {
+                    agg_kind: AggKind::Sum,
+                    return_type: DataType::Int64,
+                    inputs: vec![InputRef::new(0, DataType::Int64)],
+                    distinct: false,
+                    order_by: vec![],
+                    filter: Condition::true_cond(),
+                    direct_args: vec![],
+                };
+                let agg = Agg::new(vec![sum_agg], IndexSet::empty(), new_delete);
+                let batch_agg = BatchSimpleAgg::new(agg);
+                Ok(batch_agg.into())
+            }
+        } else {
+            let new_input = RequiredDist::single()
+                .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
+            Ok(self.clone_with_input(new_input).into())
+        }
     }
 }
 

src/frontend/src/optimizer/plan_node/batch_insert.rs

@@ -14,6 +14,8 @@
 
 use pretty_xmlish::XmlNode;
 use risingwave_common::error::Result;
+use risingwave_common::types::DataType;
+use risingwave_expr::aggregate::AggKind;
 use risingwave_pb::batch_plan::plan_node::NodeBody;
 use risingwave_pb::batch_plan::InsertNode;
 use risingwave_pb::plan_common::{DefaultColumns, IndexAndExpr};
@@ -22,10 +24,12 @@ use super::batch::prelude::*;
 use super::generic::GenericPlanRef;
 use super::utils::{childless_record, Distill};
 use super::{generic, ExprRewritable, PlanRef, PlanTreeNodeUnary, ToBatchPb, ToDistributedBatch};
-use crate::expr::Expr;
+use crate::expr::{Expr, InputRef};
 use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
-use crate::optimizer::plan_node::{PlanBase, ToLocalBatch};
+use crate::optimizer::plan_node::generic::{Agg, GenericPlanNode};
+use crate::optimizer::plan_node::{BatchSimpleAgg, PlanAggCall, PlanBase, ToLocalBatch};
 use crate::optimizer::property::{Distribution, Order, RequiredDist};
+use crate::utils::{Condition, IndexSet};
 
 /// `BatchInsert` implements [`super::LogicalInsert`]
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -36,7 +40,6 @@ pub struct BatchInsert {
 
 impl BatchInsert {
     pub fn new(core: generic::Insert<PlanRef>) -> Self {
-        assert_eq!(core.input.distribution(), &Distribution::Single);
         let base: PlanBase<Batch> =
             PlanBase::new_batch_with_core(&core, core.input.distribution().clone(), Order::any());
 
@@ -69,9 +72,43 @@ impl_plan_tree_node_for_unary! { BatchInsert }
 
 impl ToDistributedBatch for BatchInsert {
     fn to_distributed(&self) -> Result<PlanRef> {
-        let new_input = RequiredDist::single()
+        if self
+            .core
+            .ctx()
+            .session_ctx()
+            .config()
+            .batch_enable_distributed_dml()
+        {
+            // Add an hash shuffle between the insert and its input.
+            let new_input = RequiredDist::PhysicalDist(Distribution::HashShard(
+                (0..self.input().schema().len()).collect(),
+            ))
             .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
-        Ok(self.clone_with_input(new_input).into())
+            let new_insert: PlanRef = self.clone_with_input(new_input).into();
+            if self.core.returning {
+                Ok(new_insert)
+            } else {
+                let new_insert =
+                    RequiredDist::single().enforce_if_not_satisfies(new_insert, &Order::any())?;
+                // Accumulate the affected rows.
+                let sum_agg = PlanAggCall {
+                    agg_kind: AggKind::Sum,
+                    return_type: DataType::Int64,
+                    inputs: vec![InputRef::new(0, DataType::Int64)],
+                    distinct: false,
+                    order_by: vec![],
+                    filter: Condition::true_cond(),
+                    direct_args: vec![],
+                };
+                let agg = Agg::new(vec![sum_agg], IndexSet::empty(), new_insert);
+                let batch_agg = BatchSimpleAgg::new(agg);
+                Ok(batch_agg.into())
+            }
+        } else {
+            let new_input = RequiredDist::single()
+                .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
+            Ok(self.clone_with_input(new_input).into())
+        }
     }
 }
 

src/frontend/src/optimizer/plan_node/batch_update.rs

@@ -15,6 +15,8 @@
 use itertools::Itertools;
 use risingwave_common::catalog::Schema;
 use risingwave_common::error::Result;
+use risingwave_common::types::DataType;
+use risingwave_expr::aggregate::AggKind;
 use risingwave_pb::batch_plan::plan_node::NodeBody;
 use risingwave_pb::batch_plan::UpdateNode;
 
@@ -24,10 +26,12 @@ use super::utils::impl_distill_by_unit;
 use super::{
     generic, ExprRewritable, PlanBase, PlanRef, PlanTreeNodeUnary, ToBatchPb, ToDistributedBatch,
 };
-use crate::expr::{Expr, ExprRewriter, ExprVisitor};
+use crate::expr::{Expr, ExprRewriter, ExprVisitor, InputRef};
 use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
-use crate::optimizer::plan_node::ToLocalBatch;
+use crate::optimizer::plan_node::generic::{Agg, GenericPlanNode};
+use crate::optimizer::plan_node::{BatchSimpleAgg, PlanAggCall, ToLocalBatch};
 use crate::optimizer::property::{Distribution, Order, RequiredDist};
+use crate::utils::{Condition, IndexSet};
 
 /// `BatchUpdate` implements [`super::LogicalUpdate`]
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -38,9 +42,9 @@ pub struct BatchUpdate {
 
 impl BatchUpdate {
     pub fn new(core: generic::Update<PlanRef>, schema: Schema) -> Self {
-        assert_eq!(core.input.distribution(), &Distribution::Single);
         let ctx = core.input.ctx();
-        let base = PlanBase::new_batch(ctx, schema, Distribution::Single, Order::any());
+        let base =
+            PlanBase::new_batch(ctx, schema, core.input.distribution().clone(), Order::any());
         Self { base, core }
     }
 }
@@ -62,9 +66,43 @@ impl_distill_by_unit!(BatchUpdate, core, "BatchUpdate");
 
 impl ToDistributedBatch for BatchUpdate {
     fn to_distributed(&self) -> Result<PlanRef> {
-        let new_input = RequiredDist::single()
+        if self
+            .core
+            .ctx()
+            .session_ctx()
+            .config()
+            .batch_enable_distributed_dml()
+        {
+            // Add an hash shuffle between the update and its input.
+            let new_input = RequiredDist::PhysicalDist(Distribution::HashShard(
+                (0..self.input().schema().len()).collect(),
+            ))
             .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
-        Ok(self.clone_with_input(new_input).into())
+            let new_update: PlanRef = self.clone_with_input(new_input).into();
+            if self.core.returning {
+                Ok(new_update)
+            } else {
+                let new_insert =
+                    RequiredDist::single().enforce_if_not_satisfies(new_update, &Order::any())?;
+                // Accumulate the affected rows.
+                let sum_agg = PlanAggCall {
+                    agg_kind: AggKind::Sum,
+                    return_type: DataType::Int64,
+                    inputs: vec![InputRef::new(0, DataType::Int64)],
+                    distinct: false,
+                    order_by: vec![],
+                    filter: Condition::true_cond(),
+                    direct_args: vec![],
+                };
+                let agg = Agg::new(vec![sum_agg], IndexSet::empty(), new_insert);
+                let batch_agg = BatchSimpleAgg::new(agg);
+                Ok(batch_agg.into())
+            }
+        } else {
+            let new_input = RequiredDist::single()
+                .enforce_if_not_satisfies(self.input().to_distributed()?, &Order::any())?;
+            Ok(self.clone_with_input(new_input).into())
+        }
     }
 }
 

src/frontend/src/scheduler/distributed/stage.rs

@@ -708,9 +708,12 @@ impl StageRunner {
             if candidates.is_empty() {
                 return Err(SchedulerError::EmptyWorkerNodes);
             }
-            return Ok(Some(
-                candidates[self.stage.session_id.0 as usize % candidates.len()].clone(),
-            ));
+            let candidate = if self.stage.batch_enable_distributed_dml {
+                candidates[task_id as usize % candidates.len()].clone()
+            } else {
+                candidates[self.stage.session_id.0 as usize % candidates.len()].clone()
+            };
+            return Ok(Some(candidate));
         };
 
         if let Some(distributed_lookup_join_node) =