feat: support ALTER FRAGMENT SET PARALLELISM via SQL

Enable fine-grained scaling by allowing users to adjust parallelism
for individual streaming fragments using SQL. Improves operational
flexibility and user experience.

- Add Protobuf and gRPC support for fragment parallelism changes
- Extend SQL parser and AST for ALTER FRAGMENT SET PARALLELISM
- Implement frontend handler and meta client logic for new operation
- Update meta service, controller, and stream manager to reschedule
  fragments online with new parallelism
- Enhance error handling and test utilities for new trait method

Signed-off-by: Shanicky Chen <peng@risingwave-labs.com>

Author: shanicky

proto/ddl_service.proto

@@ -287,6 +287,13 @@ message AlterParallelismRequest {
 
 message AlterParallelismResponse {}
 
+message AlterFragmentParallelismRequest {
+  uint32 fragment_id = 1;
+  meta.TableParallelism parallelism = 2;
+}
+
+message AlterFragmentParallelismResponse {}
+
 message AlterCdcTableBackfillParallelismRequest {
   uint32 table_id = 1;
   meta.TableParallelism parallelism = 2;
@@ -619,6 +626,7 @@ service DdlService {
   rpc AlterOwner(AlterOwnerRequest) returns (AlterOwnerResponse);
   rpc AlterSetSchema(AlterSetSchemaRequest) returns (AlterSetSchemaResponse);
   rpc AlterParallelism(AlterParallelismRequest) returns (AlterParallelismResponse);
+  rpc AlterFragmentParallelism(AlterFragmentParallelismRequest) returns (AlterFragmentParallelismResponse);
   rpc AlterResourceGroup(AlterResourceGroupRequest) returns (AlterResourceGroupResponse);
   rpc DropTable(DropTableRequest) returns (DropTableResponse);
   rpc RisectlListStateTables(RisectlListStateTablesRequest) returns (RisectlListStateTablesResponse);

src/frontend/src/handler/alter_parallelism.rs

@@ -121,6 +121,23 @@ pub async fn handle_alter_parallelism(
     Ok(builder.into())
 }
 
+pub async fn handle_alter_fragment_parallelism(
+    handler_args: HandlerArgs,
+    fragment_id: u32,
+    parallelism: SetVariableValue,
+) -> Result<RwPgResponse> {
+    let session = handler_args.session;
+    let target_parallelism = extract_table_parallelism(parallelism)?;
+
+    session
+        .env()
+        .meta_client()
+        .alter_fragment_parallelism(fragment_id, target_parallelism)
+        .await?;
+
+    Ok(RwPgResponse::builder(StatementType::ALTER_FRAGMENT).into())
+}
+
 fn extract_table_parallelism(parallelism: SetVariableValue) -> Result<TableParallelism> {
     let adaptive_parallelism = PbTableParallelism {
         parallelism: Some(PbParallelism::Adaptive(AdaptiveParallelism {})),

src/frontend/src/handler/mod.rs

@@ -1255,17 +1255,27 @@ pub async fn handle(
         } => alter_secret::handle_alter_secret(handler_args, name, with_options, operation).await,
         Statement::AlterFragment {
             fragment_id,
-            operation: AlterFragmentOperation::AlterBackfillRateLimit { rate_limit },
-        } => {
-            alter_streaming_rate_limit::handle_alter_streaming_rate_limit_by_id(
-                &handler_args.session,
-                PbThrottleTarget::Fragment,
-                fragment_id,
-                rate_limit,
-                StatementType::SET_VARIABLE,
-            )
-            .await
-        }
+            operation,
+        } => match operation {
+            AlterFragmentOperation::AlterBackfillRateLimit { rate_limit } => {
+                alter_streaming_rate_limit::handle_alter_streaming_rate_limit_by_id(
+                    &handler_args.session,
+                    PbThrottleTarget::Fragment,
+                    fragment_id,
+                    rate_limit,
+                    StatementType::SET_VARIABLE,
+                )
+                .await
+            }
+            AlterFragmentOperation::SetParallelism { parallelism } => {
+                alter_parallelism::handle_alter_fragment_parallelism(
+                    handler_args,
+                    fragment_id,
+                    parallelism,
+                )
+                .await
+            }
+        },
         Statement::AlterDefaultPrivileges { .. } => {
             handle_privilege::handle_alter_default_privileges(handler_args, stmt).await
         }

src/frontend/src/meta_client.rs

@@ -39,7 +39,7 @@ use risingwave_pb::meta::list_streaming_job_states_response::StreamingJobState;
 use risingwave_pb::meta::list_table_fragments_response::TableFragmentInfo;
 use risingwave_pb::meta::{
     EventLog, FragmentDistribution, PbThrottleTarget, RecoveryStatus, RefreshRequest,
-    RefreshResponse,
+    RefreshResponse, TableParallelism as PbTableParallelism,
 };
 use risingwave_pb::secret::PbSecretRef;
 use risingwave_rpc_client::error::Result;
@@ -132,6 +132,12 @@ pub trait FrontendMetaClient: Send + Sync {
         rate_limit: Option<u32>,
     ) -> Result<()>;
 
+    async fn alter_fragment_parallelism(
+        &self,
+        fragment_id: u32,
+        parallelism: PbTableParallelism,
+    ) -> Result<()>;
+
     async fn get_cluster_recovery_status(&self) -> Result<RecoveryStatus>;
 
     async fn get_cluster_limits(&self) -> Result<Vec<ClusterLimit>>;
@@ -353,6 +359,16 @@ impl FrontendMetaClient for FrontendMetaClientImpl {
             .map(|_| ())
     }
 
+    async fn alter_fragment_parallelism(
+        &self,
+        fragment_id: u32,
+        parallelism: PbTableParallelism,
+    ) -> Result<()> {
+        self.0
+            .alter_fragment_parallelism(fragment_id, parallelism)
+            .await
+    }
+
     async fn get_cluster_recovery_status(&self) -> Result<RecoveryStatus> {
         self.0.get_cluster_recovery_status().await
     }

src/frontend/src/test_utils.rs

@@ -1166,6 +1166,14 @@ impl FrontendMetaClient for MockFrontendMetaClient {
         unimplemented!()
     }
 
+    async fn alter_fragment_parallelism(
+        &self,
+        _fragment_id: u32,
+        _parallelism: PbTableParallelism,
+    ) -> RpcResult<()> {
+        unimplemented!()
+    }
+
     async fn get_cluster_recovery_status(&self) -> RpcResult<RecoveryStatus> {
         Ok(RecoveryStatus::StatusRunning)
     }

src/meta/service/src/ddl_service.rs

@@ -1034,6 +1034,33 @@ impl DdlService for DdlServiceImpl {
         Ok(Response::new(AlterParallelismResponse {}))
     }
 
+    async fn alter_fragment_parallelism(
+        &self,
+        request: Request<AlterFragmentParallelismRequest>,
+    ) -> Result<Response<AlterFragmentParallelismResponse>, Status> {
+        let req = request.into_inner();
+
+        let fragment_id = req.fragment_id;
+        let parallelism = *req.get_parallelism()?;
+
+        let parallelism = match parallelism.get_parallelism()? {
+            Parallelism::Fixed(FixedParallelism { parallelism }) => {
+                StreamingParallelism::Fixed(*parallelism as _)
+            }
+            Parallelism::Auto(_) | Parallelism::Adaptive(_) => StreamingParallelism::Adaptive,
+            _ => bail_unavailable!(),
+        };
+
+        self.ddl_controller
+            .reschedule_fragment(
+                fragment_id,
+                ReschedulePolicy::Parallelism(ParallelismPolicy { parallelism }),
+            )
+            .await?;
+
+        Ok(Response::new(AlterFragmentParallelismResponse {}))
+    }
+
     /// Auto schema change for cdc sources,
     /// called by the source parser when a schema change is detected.
     async fn auto_schema_change(

src/meta/src/rpc/ddl_controller.rs

@@ -520,6 +520,17 @@ impl DdlController {
             .await
     }
 
+    pub async fn reschedule_fragment(
+        &self,
+        fragment_id: u32,
+        target: ReschedulePolicy,
+    ) -> MetaResult<()> {
+        tracing::info!("altering parallelism for fragment {}", fragment_id);
+        self.stream_manager
+            .reschedule_fragment(fragment_id, target)
+            .await
+    }
+
     async fn drop_database(&self, database_id: DatabaseId) -> MetaResult<NotificationVersion> {
         self.drop_object(ObjectType::Database, database_id as _, DropMode::Cascade)
             .await

src/meta/src/stream/scale.rs

@@ -45,6 +45,7 @@ use crate::controller::scale::{
     FragmentRenderMap, NoShuffleEnsemble, RenderedGraph, WorkerInfo,
     find_fragment_no_shuffle_dags_detailed, render_fragments, render_jobs,
 };
+use crate::error::bail_invalid_parameter;
 use crate::manager::{LocalNotification, MetaSrvEnv, MetadataManager};
 use crate::model::{ActorId, DispatcherId, FragmentId, StreamActor, StreamActorWithDispatchers};
 use crate::serving::{
@@ -459,8 +460,9 @@ impl ScaleController {
 
             let parallelism = match parallelisms.as_slice() {
                 [] => {
-                    bail!(
-                        "no reschedule policy specified for fragments in the no-shuffle ensemble: {:?}",
+                    bail_invalid_parameter!(
+                        "none of the entry fragments {:?} were included in the reschedule request; \
+                         provide at least one entry fragment id",
                         entry_fragment_ids
                     );
                 }

src/meta/src/stream/stream_manager.rs

@@ -868,6 +868,45 @@ impl GlobalStreamManager {
         Ok(())
     }
 
+    pub(crate) async fn reschedule_fragment(
+        &self,
+        fragment_id: u32,
+        target: ReschedulePolicy,
+    ) -> MetaResult<()> {
+        let _reschedule_job_lock = self.reschedule_lock_write_guard().await;
+
+        let parallelism_policy = match target {
+            ReschedulePolicy::Parallelism(policy) => policy,
+            _ => bail_invalid_parameter!("fragment reschedule only supports parallelism targets"),
+        };
+
+        let worker_nodes = self
+            .metadata_manager
+            .list_active_streaming_compute_nodes()
+            .await?
+            .into_iter()
+            .filter(|w| w.is_streaming_schedulable())
+            .collect_vec();
+        let workers = worker_nodes.into_iter().map(|x| (x.id as i32, x)).collect();
+
+        let fragment_policy = HashMap::from([(fragment_id as _, parallelism_policy.clone())]);
+
+        let commands = self
+            .scale_controller
+            .reschedule_fragment_inplace(fragment_policy, workers)
+            .await?;
+
+        let _source_pause_guard = self.source_manager.pause_tick().await;
+
+        for (database_id, command) in commands {
+            self.barrier_scheduler
+                .run_command(database_id, command)
+                .await?;
+        }
+
+        Ok(())
+    }
+
     // Don't need to add actor, just send a command
     pub async fn create_subscription(
         self: &Arc<Self>,

src/rpc_client/src/meta_client.rs

@@ -651,6 +651,20 @@ impl MetaClient {
         Ok(())
     }
 
+    pub async fn alter_fragment_parallelism(
+        &self,
+        fragment_id: u32,
+        parallelism: PbTableParallelism,
+    ) -> Result<()> {
+        let request = AlterFragmentParallelismRequest {
+            fragment_id,
+            parallelism: Some(parallelism),
+        };
+
+        self.inner.alter_fragment_parallelism(request).await?;
+        Ok(())
+    }
+
     pub async fn alter_cdc_table_backfill_parallelism(
         &self,
         table_id: u32,
@@ -2408,6 +2422,7 @@ macro_rules! for_all_meta_rpc {
             ,{ ddl_client, alter_owner, AlterOwnerRequest, AlterOwnerResponse }
             ,{ ddl_client, alter_set_schema, AlterSetSchemaRequest, AlterSetSchemaResponse }
             ,{ ddl_client, alter_parallelism, AlterParallelismRequest, AlterParallelismResponse }
+            ,{ ddl_client, alter_fragment_parallelism, AlterFragmentParallelismRequest, AlterFragmentParallelismResponse }
             ,{ ddl_client, alter_cdc_table_backfill_parallelism, AlterCdcTableBackfillParallelismRequest, AlterCdcTableBackfillParallelismResponse }
             ,{ ddl_client, alter_resource_group, AlterResourceGroupRequest, AlterResourceGroupResponse }
             ,{ ddl_client, alter_database_param, AlterDatabaseParamRequest, AlterDatabaseParamResponse }