Repository: pola-rs/polars
File: crates/polars-stream/src/graph.rs (Rust)

Code before:
use parking_lot::Mutex;
use polars_error::PolarsResult;
use slotmap::basic::Values;
use slotmap::{Key, SecondaryMap, SlotMap};

use crate::execute::StreamingExecutionState;
use crate::metrics::GraphMetrics;
use crate::nodes::ComputeNode;

slotmap::new_key_type! {
    pub struct GraphNodeKey;
    pub struct LogicalPipeKey;
}

/// Represents the compute graph.
///
/// The `nodes` perform computation and the `pipes` form the connections between nodes
/// that data is sent through.
#[derive(Default)]
pub struct Graph {
    pub nodes: SlotMap<GraphNodeKey, GraphNode>,
    pub pipes: Pipes,
}

#[derive(Clone, Default)]
pub struct Pipes(SlotMap<LogicalPipeKey, LogicalPipe>);

impl Index<LogicalPipeKey> for Pipes {
    type Output = LogicalPipe;

    fn index(&self, index: LogicalPipeKey) -> &Self::Output {
        self.0.index(index)
    }
}

impl IndexMut<LogicalPipeKey> for Pipes {
    fn index_mut(&mut self, index: LogicalPipeKey) -> &mut Self::Output {
        self.0.index_mut(index)
    }
}

impl Pipes {
    pub(crate) fn values(&self) -> Values<'_, LogicalPipeKey, LogicalPipe> {
        self.0.values()
    }
}

impl Graph {
    /// Allocate the needed `capacity` for the `Graph`.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {

Review comment (suggestion):
Is the wrapper type for this actually necessary? Feels a bit overkill/overly specific. I feel this is fine:

```rust
pub fn flush(&mut self, pipes: &SlotMap<LogicalPipeKey, LogicalPipe>) {
```

leaving the rest unchanged (besides making `pipes` public).

Code after:
use polars_error::PolarsResult;
use slotmap::{Key, SecondaryMap, SlotMap};

use crate::execute::StreamingExecutionState;
use crate::metrics::GraphMetrics;
use crate::nodes::ComputeNode;

slotmap::new_key_type! {
    pub struct GraphNodeKey;
    pub struct LogicalPipeKey;
}

/// Represents the compute graph.
///
/// The `nodes` perform computation and the `pipes` form the connections between nodes
/// that data is sent through.
#[derive(Default)]
pub struct Graph {
    pub nodes: SlotMap<GraphNodeKey, GraphNode>,
    pub pipes: SlotMap<LogicalPipeKey, LogicalPipe>,
}

impl Graph {
    /// Allocate the needed `capacity` for the `Graph`.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            nodes: SlotMap::with_capacity_and_key(capacity),
            pipes: SlotMap::with_capacity_and_key(capacity),
        }
    }

    /// Add a new `GraphNode` to the `Graph` and connect the inputs and outputs
    /// to their respective `LogicalPipe`s.
    pub fn add_node<N: ComputeNode + 'static>(
        &mut self,
        node: N,
        inputs: impl IntoIterator<Item = (GraphNodeKey, usize)>,
    ) -> GraphNodeKey {
        // Add the GraphNode.
        let node_key = self.nodes.insert(GraphNode {
            compute: Box::new(node),
            inputs: Vec::new(),
            outputs: Vec::new(),
        });

        // Create and add pipes that connect input to output.
        for (recv_port, (sender, send_port)) in inputs.into_iter().enumerate() {
            let pipe = LogicalPipe {
                sender,
                send_port,
                send_state: PortState::Blocked,