refactor(query): use graphs to build pipelines (#18184)

* refactor(query): use graphs to build pipelines

* refactor(query): use graphs to build pipelines

* refactor(query): use graphs to build pipelines

* refactor(query): use graphs to build pipelines

* refactor(query): use graphs to build pipelines

* refactor(query): use graphs to build pipelines

Author: zhang2014

src/query/pipeline/core/src/pipeline.rs

@@ -15,7 +15,8 @@
 use std::fmt::Display;
 use std::fmt::Formatter;
 
-use crate::pipe::Pipe;
+use petgraph::dot::Dot;
+
 use crate::Pipeline;
 
 impl Pipeline {
@@ -28,35 +29,17 @@ struct PipelineIndentDisplayWrapper<'a> {
     pipeline: &'a Pipeline,
 }
 
-impl PipelineIndentDisplayWrapper<'_> {
-    fn pipe_name(pipe: &Pipe) -> String {
-        unsafe { pipe.items[0].processor.name() }
-    }
-}
-
 impl Display for PipelineIndentDisplayWrapper<'_> {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        let pipes = &self.pipeline.pipes;
-        for (index, pipe) in pipes.iter().rev().enumerate() {
-            if index > 0 {
-                writeln!(f)?;
-            }
-
-            for _ in 0..index {
-                write!(f, "  ")?;
-            }
-
-            let pipe_name = Self::pipe_name(pipe);
-            if pipe.input_length == pipe.output_length
-                || pipe.input_length == 0
-                || pipe.output_length == 0
-            {
-                write!(f, "{} × {}", Self::pipe_name(pipe), pipe.items.len(),)?;
-            } else {
-                write!(f, "Merge to {pipe_name} × {}", pipe.output_length,)?;
-            }
-        }
-
-        Ok(())
+        write!(
+            f,
+            "{:?}",
+            Dot::with_attr_getters(
+                &self.pipeline.graph,
+                &[],
+                &|_, _| String::new(),
+                &|_, (_, _)| String::new(),
+            )
+        )
     }
 }

src/query/pipeline/core/src/pipeline_display.rs

@@ -12,13 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cell::RefCell;
 use std::collections::btree_map::Entry;
 use std::collections::BTreeMap;
-use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
-use databend_common_base::runtime::drop_guard;
 use databend_common_base::runtime::error_info::NodeErrorType;
 use databend_common_base::runtime::metrics::MetricSample;
 use databend_common_base::runtime::metrics::ScopedRegistry;
@@ -29,25 +28,18 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::StackTrace;
 
 pub struct PlanScopeGuard {
-    idx: usize,
-    scope_size: Arc<AtomicUsize>,
+    save: Option<Arc<PlanScope>>,
 }
 
 impl PlanScopeGuard {
-    pub fn create(scope_size: Arc<AtomicUsize>, idx: usize) -> PlanScopeGuard {
-        PlanScopeGuard { idx, scope_size }
+    pub fn new(save: Option<Arc<PlanScope>>) -> PlanScopeGuard {
+        PlanScopeGuard { save }
     }
 }
 
 impl Drop for PlanScopeGuard {
     fn drop(&mut self) {
-        drop_guard(move || {
-            if self.scope_size.fetch_sub(1, Ordering::SeqCst) != self.idx + 1
-                && !std::thread::panicking()
-            {
-                panic!("Broken pipeline scope stack.");
-            }
-        })
+        PLAN_SCOPE.replace(self.save.take());
     }
 }
 
@@ -186,20 +178,34 @@ pub struct PlanScope {
     pub metrics_registry: Arc<ScopedRegistry>,
 }
 
+thread_local! {
+    static PLAN_SCOPE: RefCell<Option<Arc<PlanScope>>> = const { RefCell::new(None) };
+}
+
 impl PlanScope {
     pub fn create(
         id: u32,
         name: String,
         title: Arc<String>,
         labels: Arc<Vec<ProfileLabel>>,
-    ) -> PlanScope {
-        PlanScope {
+    ) -> Arc<PlanScope> {
+        let parent_id = PlanScope::get_plan_scope().map(|x| x.id);
+        Arc::new(PlanScope {
             id,
             labels,
             title,
-            parent_id: None,
+            parent_id,
             name,
             metrics_registry: ScopedRegistry::create(None),
-        }
+        })
+    }
+
+    pub fn enter_scope_guard(self: Arc<Self>) -> PlanScopeGuard {
+        let old = PLAN_SCOPE.replace(Some(self));
+        PlanScopeGuard::new(old)
+    }
+
+    pub fn get_plan_scope() -> Option<Arc<PlanScope>> {
+        PLAN_SCOPE.with(|x| x.borrow().clone())
     }
 }

src/query/pipeline/core/src/processors/profile.rs

@@ -36,6 +36,7 @@
 #![feature(try_blocks)]
 #![feature(variant_count)]
 #![feature(duration_constructors)]
+#![feature(get_mut_unchecked)]
 #![allow(clippy::diverging_sub_expression)]
 #![allow(clippy::arc_with_non_send_sync)]
 

src/query/service/src/lib.rs

@@ -12,7 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
+
 use databend_common_exception::Result;
+use databend_common_pipeline_core::PlanScope;
 use databend_common_sql::executor::physical_plans::ExchangeSink;
 use databend_common_sql::executor::physical_plans::ExchangeSource;
 
@@ -27,8 +30,21 @@ impl PipelineBuilder {
             self.exchange_injector.clone(),
         )?;
 
-        // add profile
-        build_res.main_pipeline.reset_scopes(&self.main_pipeline);
+        let plan_scope = PlanScope::get_plan_scope();
+        for node in build_res.main_pipeline.graph.node_weights_mut() {
+            let Some(scope) = node.scope.as_mut() else {
+                node.scope = plan_scope.clone();
+                continue;
+            };
+
+            if scope.parent_id.is_none() {
+                unsafe {
+                    let scope = Arc::get_mut_unchecked(scope);
+                    scope.parent_id = plan_scope.as_ref().map(|x| x.id);
+                }
+            }
+        }
+
         // add sharing data
         self.join_state = build_res.builder_data.input_join_state;
         self.merge_into_probe_data_fields = build_res.builder_data.input_probe_schema;

src/query/service/src/pipelines/builders/builder_exchange.rs

@@ -38,12 +38,8 @@ impl PipelineBuilder {
     // Create a new pipeline builder with the same context as the current builder
     fn create_sub_pipeline_builder(&self) -> PipelineBuilder {
         let sub_context = QueryContext::create_from(self.ctx.as_ref());
-        let mut sub_builder = PipelineBuilder::create(
-            self.func_ctx.clone(),
-            self.settings.clone(),
-            sub_context,
-            self.main_pipeline.get_scopes(),
-        );
+        let mut sub_builder =
+            PipelineBuilder::create(self.func_ctx.clone(), self.settings.clone(), sub_context);
         sub_builder.hash_join_states = self.hash_join_states.clone();
         sub_builder
     }

src/query/service/src/pipelines/builders/builder_join.rs

@@ -72,12 +72,8 @@ impl PipelineBuilder {
 
     fn expand_union_all(&mut self, input: &PhysicalPlan) -> Result<Receiver<DataBlock>> {
         let union_ctx = QueryContext::create_from(self.ctx.as_ref());
-        let mut pipeline_builder = PipelineBuilder::create(
-            self.func_ctx.clone(),
-            self.settings.clone(),
-            union_ctx,
-            self.main_pipeline.get_scopes(),
-        );
+        let mut pipeline_builder =
+            PipelineBuilder::create(self.func_ctx.clone(), self.settings.clone(), union_ctx);
         pipeline_builder.hash_join_states = self.hash_join_states.clone();
 
         let mut build_res = pipeline_builder.finalize(input)?;

src/query/service/src/pipelines/builders/builder_union_all.rs

@@ -96,7 +96,7 @@ pub(crate) struct Node {
 impl Node {
     pub fn create(
         pid: usize,
-        scope: Option<PlanScope>,
+        scope: Option<Arc<PlanScope>>,
         processor: &ProcessorPtr,
         inputs_port: &[Arc<InputPort>],
         outputs_port: &[Arc<OutputPort>],
@@ -237,68 +237,32 @@ impl ExecutingGraph {
         graph: &mut StableGraph<Arc<Node>, EdgeInfo>,
         time_series_profile_builder: &mut QueryTimeSeriesProfileBuilder,
     ) {
-        #[derive(Debug)]
-        struct Edge {
-            source_port: usize,
-            source_node: NodeIndex,
-            target_port: usize,
-            target_node: NodeIndex,
-        }
-
-        let mut pipes_edges: Vec<Vec<Edge>> = Vec::new();
-        for pipe in &pipeline.pipes {
-            assert_eq!(
-                pipe.input_length,
-                pipes_edges.last().map(|x| x.len()).unwrap_or_default()
-            );
-
-            let mut edge_index = 0;
-            let mut pipe_edges = Vec::with_capacity(pipe.output_length);
-
-            for item in &pipe.items {
-                let pid = graph.node_count();
-                let time_series_profile = if let Some(scope) = pipe.scope.as_ref() {
-                    let plan_id = scope.id;
-                    let time_series_profile =
-                        time_series_profile_builder.register_time_series_profile(plan_id);
-                    Some(time_series_profile)
-                } else {
-                    None
-                };
-                let node = Node::create(
-                    pid,
-                    pipe.scope.clone(),
-                    &item.processor,
-                    &item.inputs_port,
-                    &item.outputs_port,
-                    time_series_profile,
-                );
-
-                let graph_node_index = graph.add_node(node.clone());
-                unsafe {
-                    item.processor.set_id(graph_node_index);
-                }
-
-                for offset in 0..item.inputs_port.len() {
-                    let last_edges = pipes_edges.last_mut().unwrap();
+        let offset = graph.node_count();
+        for node in pipeline.graph.node_weights() {
+            let pid = graph.node_count();
+            let mut time_series_profile = None;
+
+            if let Some(scope) = node.scope.as_ref() {
+                let plan_id = scope.id;
+                time_series_profile =
+                    Some(time_series_profile_builder.register_time_series_profile(plan_id));
+            }
 
-                    last_edges[edge_index].target_port = offset;
-                    last_edges[edge_index].target_node = graph_node_index;
-                    edge_index += 1;
-                }
+            let graph_node_index = graph.add_node(Node::create(
+                pid,
+                node.scope.clone(),
+                &node.proc,
+                &node.inputs,
+                &node.outputs,
+                time_series_profile,
+            ));
 
-                for offset in 0..item.outputs_port.len() {
-                    pipe_edges.push(Edge {
-                        source_port: offset,
-                        source_node: graph_node_index,
-                        target_port: 0,
-                        target_node: Default::default(),
-                    });
-                }
+            unsafe {
+                node.proc.set_id(graph_node_index);
             }
-
-            pipes_edges.push(pipe_edges);
         }
+
+        // FIXME:
         let query_time_series = Arc::new(time_series_profile_builder.build());
         let node_indices: Vec<_> = graph.node_indices().collect();
         for node_index in node_indices {
@@ -313,23 +277,24 @@ impl ExecutingGraph {
             }
         }
 
-        // The last pipe cannot contain any output edge.
-        assert!(pipes_edges.last().map(|x| x.is_empty()).unwrap_or_default());
-        pipes_edges.pop();
+        for edge in pipeline.graph.edge_indices() {
+            let index = EdgeIndex::new(edge.index());
+            if let Some((source, target)) = pipeline.graph.edge_endpoints(index) {
+                let source = NodeIndex::new(offset + source.index());
+                let target = NodeIndex::new(offset + target.index());
+                let edge_weight = pipeline.graph.edge_weight(index).unwrap();
 
-        for pipe_edges in &pipes_edges {
-            for edge in pipe_edges {
-                let edge_index = graph.add_edge(edge.source_node, edge.target_node, EdgeInfo {
-                    input_index: edge.target_port,
-                    output_index: edge.source_port,
+                let edge_index = graph.add_edge(source, target, EdgeInfo {
+                    input_index: edge_weight.input_index,
+                    output_index: edge_weight.output_index,
                 });
 
                 unsafe {
-                    let (target_node, target_port) = (edge.target_node, edge.target_port);
+                    let (target_node, target_port) = (target, edge_weight.input_index);
                     let input_trigger = graph[target_node].create_trigger(edge_index);
                     graph[target_node].inputs_port[target_port].set_trigger(input_trigger);
 
-                    let (source_node, source_port) = (edge.source_node, edge.source_port);
+                    let (source_node, source_port) = (source, edge_weight.output_index);
                     let output_trigger = graph[source_node].create_trigger(edge_index);
                     graph[source_node].outputs_port[source_port].set_trigger(output_trigger);
 

src/query/service/src/pipelines/executor/executor_graph.rs

@@ -22,7 +22,6 @@ mod global_queries_executor;
 mod pipeline_complete_executor;
 mod pipeline_executor;
 mod pipeline_pulling_executor;
-mod pipeline_pushing_executor;
 mod processor_async_task;
 mod queries_executor_tasks;
 mod queries_pipeline_executor;
@@ -40,7 +39,6 @@ pub use global_queries_executor::GlobalQueriesExecutor;
 pub use pipeline_complete_executor::PipelineCompleteExecutor;
 pub use pipeline_executor::PipelineExecutor;
 pub use pipeline_pulling_executor::PipelinePullingExecutor;
-pub use pipeline_pushing_executor::PipelinePushingExecutor;
 pub use processor_async_task::ProcessorAsyncTask;
 pub use queries_executor_tasks::QueriesExecutorTasksQueue;
 pub use queries_pipeline_executor::QueriesPipelineExecutor;

src/query/service/src/pipelines/executor/mod.rs

@@ -102,7 +102,7 @@ pub struct PipelinePullingExecutor {
 
 impl PipelinePullingExecutor {
     fn wrap_pipeline(pipeline: &mut Pipeline, tx: SyncSender<DataBlock>) -> Result<()> {
-        if pipeline.is_pushing_pipeline()? || !pipeline.is_pulling_pipeline()? {
+        if !pipeline.is_pulling_pipeline()? {
             return Err(ErrorCode::Internal(
                 "Logical error, PipelinePullingExecutor can only work on pulling pipeline.",
             ));