update

luoluoyuyu · luoluoyuyu · commit 4eb4d533a96c · 2026-04-19T18:53:08.000+08:00
diff --git a/conf/config.yaml b/conf/config.yaml
@@ -51,15 +51,17 @@ wasm:
 
 # Streaming Runtime Configuration
 streaming:
-  # Bytes in the global memory pool for streaming execution: pipeline buffers, batch collect,
-  # backpressure. Omitted → 200 MiB.
-  streaming_runtime_memory_bytes: 209715200
+  # Global memory pool for streaming pipeline execution (buffers, batch collect, backpressure).
+  # Default / example: 10 MiB (10485760 bytes).
+  streaming_runtime_memory_bytes: 10485760
 
-  # Per stateful operator (join / agg / window): in-memory state store cap; spill when exceeded.
-  # Omitted → 100 MiB.
-  operator_state_store_memory_bytes: 104857600
+  # Per stateful operator (join / agg / window): in-memory state store cap before spill.
+  # Default / example: 5 MiB (5242880 bytes).
+  operator_state_store_memory_bytes: 5242880
   checkpoint_interval_ms: 60000
   pipeline_parallelism: 1
+  # KeyBy (key extraction) operator pipeline parallelism in planned streaming jobs.
+  key_by_parallelism: 1
 
 # State Storage Configuration
 # Used to store runtime state data for tasks
diff --git a/src/config/global_config.rs b/src/config/global_config.rs
@@ -20,21 +20,21 @@ use crate::config::service_config::ServiceConfig;
 use crate::config::streaming_job::{ResolvedStreamingJobConfig, StreamingJobConfig};
 use crate::config::wasm_config::WasmConfig;
 
-/// Default for [`StreamingConfig::streaming_runtime_memory_bytes`] when unset. **200 MiB.**
-pub const DEFAULT_STREAMING_RUNTIME_MEMORY_BYTES: u64 = 200 * 1024 * 1024;
+/// Default for [`StreamingConfig::streaming_runtime_memory_bytes`] when unset. **10 MiB** (pipeline buffers, backpressure).
+pub const DEFAULT_STREAMING_RUNTIME_MEMORY_BYTES: u64 = 10 * 1024 * 1024;
 
-/// Default for [`StreamingConfig::operator_state_store_memory_bytes`] when unset. **100 MiB.**
-pub const DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES: u64 = 100 * 1024 * 1024;
+/// Default for [`StreamingConfig::operator_state_store_memory_bytes`] when unset. **5 MiB** per stateful operator cap.
+pub const DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES: u64 = 5 * 1024 * 1024;
 
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct StreamingConfig {
     #[serde(flatten)]
     pub job: StreamingJobConfig,
-    /// Bytes reserved in the global memory pool for streaming execution (pipeline buffers,
-    /// batch collect, backpressure).
+    /// Bytes reserved in the global memory pool for streaming pipeline execution (buffers,
+    /// batch collect, backpressure). Default 10 MiB.
     #[serde(default)]
     pub streaming_runtime_memory_bytes: Option<u64>,
-    /// Per stateful operator: in-memory state store cap before spill.
+    /// Per stateful operator: in-memory state store cap before spill. Default 5 MiB.
     #[serde(default)]
     pub operator_state_store_memory_bytes: Option<u64>,
 }
diff --git a/src/config/streaming_job.rs b/src/config/streaming_job.rs
@@ -14,19 +14,24 @@ use serde::{Deserialize, Serialize};
 
 pub const DEFAULT_CHECKPOINT_INTERVAL_MS: u64 = 60 * 1000;
 pub const DEFAULT_PIPELINE_PARALLELISM: u32 = 1;
+pub const DEFAULT_KEY_BY_PARALLELISM: u32 = 1;
 
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct StreamingJobConfig {
     #[serde(default)]
     pub checkpoint_interval_ms: Option<u64>,
     #[serde(default)]
     pub pipeline_parallelism: Option<u32>,
+    /// Physical parallelism for KeyBy / key-extraction operators in planned streaming graphs.
+    #[serde(default)]
+    pub key_by_parallelism: Option<u32>,
 }
 
 #[derive(Debug, Clone, Copy)]
 pub struct ResolvedStreamingJobConfig {
     pub checkpoint_interval_ms: u64,
     pub pipeline_parallelism: u32,
+    pub key_by_parallelism: u32,
 }
 
 impl StreamingJobConfig {
@@ -40,6 +45,10 @@ impl StreamingJobConfig {
                 .pipeline_parallelism
                 .filter(|&p| p > 0)
                 .unwrap_or(DEFAULT_PIPELINE_PARALLELISM),
+            key_by_parallelism: self
+                .key_by_parallelism
+                .filter(|&p| p > 0)
+                .unwrap_or(DEFAULT_KEY_BY_PARALLELISM),
         }
     }
 }
diff --git a/src/coordinator/execution/executor.rs b/src/coordinator/execution/executor.rs
@@ -323,11 +323,16 @@ impl PlanVisitor for Executor {
         let execute = || -> Result<ExecuteResult, ExecuteError> {
             let mut fs_program: FsProgram = plan.program.clone().into();
             let job_manager: Arc<JobManager> = Arc::clone(&self.job_manager);
-            let pipeline_parallelism = parse_pipeline_parallelism(plan.with_options.as_ref())
-                .unwrap_or_else(|| job_manager.default_pipeline_parallelism())
-                .max(1);
-            for node in &mut fs_program.nodes {
-                node.parallelism = pipeline_parallelism;
+            // Only override per-node parallelism when CREATE STREAMING TABLE specifies
+            // `WITH (parallelism = N)`. Otherwise keep planner-assigned values (e.g. keyed
+            // aggregates defaulting to a higher parallelism than the job-wide default).
+            if let Some(pipeline_parallelism) =
+                parse_pipeline_parallelism(plan.with_options.as_ref())
+            {
+                let p = pipeline_parallelism.max(1);
+                for node in &mut fs_program.nodes {
+                    node.parallelism = p;
+                }
             }
 
             let job_id = plan.name.clone();
diff --git a/src/server/initializer.rs b/src/server/initializer.rs
@@ -19,6 +19,12 @@ use crate::config::GlobalConfig;
 
 pub type InitializerFn = fn(&GlobalConfig) -> Result<()>;
 
+fn initialize_streaming_sql_planning(config: &GlobalConfig) -> Result<()> {
+    let job = config.streaming.resolved_job();
+    crate::sql::planning_runtime::install_sql_planning_from_streaming_job(&job);
+    Ok(())
+}
+
 #[derive(Clone)]
 pub struct Component {
     pub name: &'static str,
@@ -94,6 +100,7 @@ impl ComponentRegistry {
 pub fn build_core_registry() -> ComponentRegistry {
     let builder = {
         let b = ComponentRegistryBuilder::new()
+            .register("StreamingSqlPlanning", initialize_streaming_sql_planning)
             .register("WasmCache", initialize_wasm_cache)
             .register("TaskManager", initialize_task_manager)
             .register("MemoryService", initialize_memory_service)
diff --git a/src/sql/common/constants.rs b/src/sql/common/constants.rs
@@ -108,6 +108,8 @@ pub mod sql_field {
 
 pub mod sql_planning_default {
     pub const DEFAULT_PARALLELISM: usize = 1;
+    /// Default physical parallelism for `KeyBy` / key-extraction pipelines (configurable via YAML).
+    pub const DEFAULT_KEY_BY_PARALLELISM: usize = 1;
     /// Parallelism for aggregations that run after `KeyBy` / shuffle on non-empty routing keys.
     pub const KEYED_AGGREGATE_DEFAULT_PARALLELISM: usize = 8;
     pub const PLANNING_TTL_SECS: u64 = 24 * 60 * 60;
diff --git a/src/sql/logical_node/aggregate.rs b/src/sql/logical_node/aggregate.rs
@@ -64,12 +64,11 @@ multifield_partial_ord!(
 );
 
 impl StreamWindowAggregateNode {
+    /// This node is only emitted after `KeyExtractionNode` in streaming rewrites; `partition_keys`
+    /// may be empty when GROUP BY is only a window call (window column stripped from key list),
+    /// but the pipeline still consumes a shuffle — use keyed aggregate parallelism.
     fn parallelism_after_keyed_shuffle(&self, planner: &Planner) -> usize {
-        if self.partition_keys.is_empty() {
-            planner.default_parallelism()
-        } else {
-            planner.keyed_aggregate_parallelism()
-        }
+        planner.keyed_aggregate_parallelism()
     }
 
     /// Safely constructs a new node, computing the final projection without panicking.
diff --git a/src/sql/logical_node/key_calculation.rs b/src/sql/logical_node/key_calculation.rs
@@ -238,7 +238,7 @@ impl StreamingOperatorBlueprint for KeyExtractionNode {
             engine_operator_name,
             protobuf_payload,
             format!("Key<{}>", self.operator_label.as_deref().unwrap_or("_")),
-            planner.default_parallelism(),
+            planner.key_by_parallelism(),
         );
 
         let data_edge =
diff --git a/src/sql/logical_node/logical/operator_chain.rs b/src/sql/logical_node/logical/operator_chain.rs
@@ -128,4 +128,17 @@ impl OperatorChain {
     pub fn is_sink(&self) -> bool {
         self.operators[0].operator_name == OperatorName::ConnectorSink
     }
+
+    /// Operators safe to run at a higher upstream `TaskContext::parallelism` when fused after a
+    /// stateful node (e.g. window aggregate @ 8 → projection @ 1).
+    pub fn is_parallelism_upstream_expandable(&self) -> bool {
+        self.operators.iter().all(|op| {
+            matches!(
+                op.operator_name,
+                OperatorName::Projection
+                    | OperatorName::Value
+                    | OperatorName::ExpressionWatermark
+            )
+        })
+    }
 }
diff --git a/src/sql/logical_node/updating_aggregate.rs b/src/sql/logical_node/updating_aggregate.rs
@@ -218,11 +218,7 @@ impl StreamingOperatorBlueprint for ContinuousAggregateNode {
 
         let operator_config = self.compile_operator_config(planner, &upstream_schema)?;
 
-        let parallelism = if self.partition_key_indices.is_empty() {
-            planner.default_parallelism()
-        } else {
-            planner.keyed_aggregate_parallelism()
-        };
+        let parallelism = planner.keyed_aggregate_parallelism();
 
         let logical_node = LogicalNode::single(
             node_index as u32,
diff --git a/src/sql/logical_node/windows_function.rs b/src/sql/logical_node/windows_function.rs
@@ -163,11 +163,7 @@ impl StreamingOperatorBlueprint for StreamingWindowFunctionNode {
             window_function_plan: evaluation_plan_payload,
         };
 
-        let parallelism = if self.partition_key_indices.is_empty() {
-            planner.default_parallelism()
-        } else {
-            planner.keyed_aggregate_parallelism()
-        };
+        let parallelism = planner.keyed_aggregate_parallelism();
 
         let logical_node = LogicalNode::single(
             node_index as u32,
diff --git a/src/sql/logical_planner/optimizers/chaining.rs b/src/sql/logical_planner/optimizers/chaining.rs
@@ -45,9 +45,14 @@ impl Optimizer for ChainingOptimizer {
             let source_node = plan.node_weight(node_idx).expect("Source node missing");
             let target_node = plan.node_weight(target_idx).expect("Target node missing");
 
+            let parallelism_ok = source_node.parallelism == target_node.parallelism
+                || target_node
+                    .operator_chain
+                    .is_parallelism_upstream_expandable();
+
             if source_node.operator_chain.is_source()
                 || target_node.operator_chain.is_sink()
-                || source_node.parallelism != target_node.parallelism
+                || !parallelism_ok
             {
                 continue;
             }
@@ -93,6 +98,10 @@ impl Optimizer for ChainingOptimizer {
             source_node.description =
                 format!("{} -> {}", source_node.description, target_node.description);
 
+            source_node.parallelism = source_node
+                .parallelism
+                .max(target_node.parallelism);
+
             source_node
                 .operator_chain
                 .operators
@@ -150,6 +159,28 @@ mod tests {
         )
     }
 
+    /// Window aggregate at higher default parallelism may forward into projection @ 1: still fuse
+    /// so each branch does not reserve a separate global state-memory block for the same sub-chain.
+    #[test]
+    fn fusion_stateful_high_parallelism_into_expandable_low() {
+        let mut g = LogicalGraph::new();
+        let n0 = g.add_node(source_node());
+        let n1 = g.add_node(proj_node(1, "tumble"));
+        let n2 = g.add_node(proj_node(2, "proj"));
+        let n1w = g.node_weight_mut(n1).unwrap();
+        n1w.parallelism = 8;
+        let e = forward_edge();
+        g.add_edge(n0, n1, e.clone());
+        g.add_edge(n1, n2, e);
+
+        let changed = ChainingOptimizer {}.optimize_once(&mut g);
+        assert!(changed);
+        assert_eq!(g.node_count(), 2);
+        let fused = g.node_weights().find(|n| n.description.contains("->")).unwrap();
+        assert_eq!(fused.parallelism, 8);
+        assert_eq!(fused.operator_chain.len(), 2);
+    }
+
     /// Regression: upstream at last `NodeIndex` + remove non-last downstream swaps indices.
     #[test]
     fn fusion_remaps_when_upstream_was_last_node_index() {
diff --git a/src/sql/logical_planner/streaming_planner.rs b/src/sql/logical_planner/streaming_planner.rs
@@ -102,6 +102,11 @@ impl<'a> Planner<'a> {
         self.schema_provider.default_parallelism()
     }
 
+    #[inline]
+    pub(crate) fn key_by_parallelism(&self) -> usize {
+        self.schema_provider.key_by_parallelism()
+    }
+
     /// Parallelism for operators that consume a keyed shuffle (non-empty partition keys).
     #[inline]
     pub(crate) fn keyed_aggregate_parallelism(&self) -> usize {
diff --git a/src/sql/mod.rs b/src/sql/mod.rs
@@ -19,6 +19,7 @@ pub mod logical_node;
 pub mod logical_planner;
 pub mod parse;
 pub mod physical;
+pub(crate) mod planning_runtime;
 pub mod schema;
 pub mod types;
 
diff --git a/src/sql/schema/schema_provider.rs b/src/sql/schema/schema_provider.rs
@@ -145,14 +145,21 @@ impl StreamPlanningContext {
         self.sql_config.default_parallelism
     }
 
+    #[inline]
+    pub fn key_by_parallelism(&self) -> usize {
+        self.sql_config.key_by_parallelism
+    }
+
     /// Same registration order as the historical `StreamSchemaProvider::new` (placeholders, then DataFusion defaults).
     pub fn new() -> Self {
-        Self::builder()
+        let mut ctx = Self::builder()
             .with_streaming_extensions()
             .expect("streaming extensions")
             .with_default_functions()
             .expect("default functions")
-            .build()
+            .build();
+        ctx.sql_config = crate::sql::planning_runtime::sql_planning_snapshot();
+        ctx
     }
 
     pub fn register_stream_table(&mut self, table: StreamTable) {
diff --git a/src/sql/types/mod.rs b/src/sql/types/mod.rs
@@ -38,12 +38,15 @@ pub enum ProcessingMode {
 #[derive(Clone, Debug)]
 pub struct SqlConfig {
     pub default_parallelism: usize,
+    /// Physical pipeline parallelism for [`KeyExtractionNode`](crate::sql::logical_node::key_calculation::KeyExtractionNode) / KeyBy.
+    pub key_by_parallelism: usize,
 }
 
 impl Default for SqlConfig {
     fn default() -> Self {
         Self {
             default_parallelism: sql_planning_default::DEFAULT_PARALLELISM,
+            key_by_parallelism: sql_planning_default::DEFAULT_KEY_BY_PARALLELISM,
         }
     }
 }
