feat: Highlight nodes in streaming phys plan graph (#22535)

Author: JakubValtar

crates/polars-stream/src/physical_plan/fmt.rs

@@ -12,6 +12,62 @@ use slotmap::{Key, SecondaryMap, SlotMap};
 
 use super::{PhysNode, PhysNodeKey, PhysNodeKind};
 
+/// A style of a graph node.
+enum NodeStyle {
+    InMemoryFallback,
+    MemoryIntensive,
+    Generic,
+}
+
+impl NodeStyle {
+    const COLOR_IN_MEM_FALLBACK: &str = "0.0 0.3 1.0"; // Pastel red
+    const COLOR_MEM_INTENSIVE: &str = "0.16 0.3 1.0"; // Pastel yellow
+
+    /// Returns a style for a node kind.
+    pub fn for_node_kind(kind: &PhysNodeKind) -> Self {
+        use PhysNodeKind as K;
+        match kind {
+            K::InMemoryMap { .. } => Self::InMemoryFallback,
+            K::InMemorySource { .. }
+            | K::InputIndependentSelect { .. }
+            | K::NegativeSlice { .. }
+            | K::InMemorySink { .. }
+            | K::Sort { .. }
+            | K::GroupBy { .. }
+            | K::EquiJoin { .. }
+            | K::SemiAntiJoin { .. }
+            | K::InMemoryJoin { .. }
+            | K::MergeSorted { .. }
+            | K::Multiplexer { .. } => Self::MemoryIntensive,
+            _ => Self::Generic,
+        }
+    }
+
+    /// Returns extra styling attributes (if any) for the graph node.
+    pub fn node_attrs(&self) -> Option<String> {
+        match self {
+            Self::InMemoryFallback => Some(format!(
+                "style=filled,fillcolor=\"{}\"",
+                Self::COLOR_IN_MEM_FALLBACK
+            )),
+            Self::MemoryIntensive => Some(format!(
+                "style=filled,fillcolor=\"{}\"",
+                Self::COLOR_MEM_INTENSIVE
+            )),
+            Self::Generic => None,
+        }
+    }
+
+    /// Returns a legend explaining the node style meaning.
+    pub fn legend() -> String {
+        format!(
+            "fontsize=\"10\"\nlabelloc=\"b\"\nlabel=<<BR/><BR/><B>Legend</B><BR/><BR/>◯ streaming engine node<FONT COLOR=\"{}\">⬤</FONT>potentially memory-intensive node<FONT COLOR=\"{}\">⬤</FONT>in-memory engine fallback>",
+            Self::COLOR_MEM_INTENSIVE,
+            Self::COLOR_IN_MEM_FALLBACK,
+        )
+    }
+}
+
 fn escape_graphviz(s: &str) -> String {
     s.replace('\\', "\\\\")
         .replace('\n', "\\n")
@@ -39,8 +95,10 @@ fn visualize_plan_rec(
     }
     visited.insert(node_key, ());
 
+    let kind = &phys_sm[node_key].kind;
+
     use std::slice::from_ref;
-    let (label, inputs) = match &phys_sm[node_key].kind {
+    let (label, inputs) = match kind {
         PhysNodeKind::InMemorySource { df } => (
             format!(
                 "in-memory-source\\ncols: {}",
@@ -329,11 +387,14 @@ fn visualize_plan_rec(
         },
     };
 
-    out.push(format!(
-        "{} [label=\"{}\"];",
-        node_key.data().as_ffi(),
-        label
-    ));
+    let node_id = node_key.data().as_ffi();
+    let style = NodeStyle::for_node_kind(kind);
+
+    if let Some(attrs) = style.node_attrs() {
+        out.push(format!("{node_id} [label=\"{label}\",{attrs}];"));
+    } else {
+        out.push(format!("{node_id} [label=\"{label}\"];"));
+    }
     for input in inputs {
         visualize_plan_rec(input.node, phys_sm, expr_arena, visited, out);
         out.push(format!(
@@ -350,8 +411,9 @@ pub fn visualize_plan(
     expr_arena: &Arena<AExpr>,
 ) -> String {
     let mut visited: SecondaryMap<PhysNodeKey, ()> = SecondaryMap::new();
-    let mut out = Vec::with_capacity(phys_sm.len() + 2);
+    let mut out = Vec::with_capacity(phys_sm.len() + 3);
     out.push("digraph polars {\nrankdir=\"BT\"".to_string());
+    out.push(NodeStyle::legend());
     visualize_plan_rec(root, phys_sm, expr_arena, &mut visited, &mut out);
     out.push("}".to_string());
     out.join("\n")