07-polars Streaming

        match engine {
            Engine::Streaming => {
                feature_gated!("streaming", self = self.with_streaming(true))
            },
            Engine::Gpu => self = self.with_gpu(true),
            _ => (),
        }

        let mut ir_plan = self.to_alp_optimized()?;

        ir_plan.ensure_root_node_is_sink();

        match engine {
            Engine::Streaming => feature_gated!("streaming", {
                polars_stream::run_query(
                    ir_plan.lp_top,
                    &mut ir_plan.lp_arena,
                    &mut ir_plan.expr_arena,
                )
            }),

/// Executes the IR with the streaming engine.
///
/// Unsupported operations can fall back to the in-memory engine.
///
/// Returns:
/// - `Ok(QueryResult::Single(DataFrame))` when collecting to a single sink.
/// - `Ok(QueryResult::Multiple(Vec<DataFrame>))` when collecting to multiple sinks.
/// - `Err` if the IR can't be executed.
///
/// Returned `DataFrame`s contain data only for memory sinks,
/// `DataFrame`s corresponding to file sinks are empty.
pub fn run_query(
    node: Node,
    ir_arena: &mut Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<QueryResult> {
    StreamingQuery::build(node, ir_arena, expr_arena)?.execute()
}

#[derive(Debug)]
pub struct Morsel {
    /// The data contained in this morsel.
    sf: SpillFrame,

    /// The sequence number of this morsel. May only stay equal or increase
    /// within a pipeline.
    seq: MorselSeq,

    /// A token that indicates which source this morsel originates from.
    source_token: SourceToken,

    /// Used to notify someone when this morsel is consumed, to provide backpressure.
    consume_token: Option<WaitToken>,
}

impl StreamingQuery {
    pub fn build(
        node: Node,
        ir_arena: &mut Arena<IR>,
        expr_arena: &mut Arena<AExpr>,
    ) -> PolarsResult<Self> {

        let root_phys_node = crate::physical_plan::build_physical_plan(
            node,
            ir_arena,
            expr_arena,
            &mut phys_sm,
            ctx,
        )?;
        if let Ok(visual_path) = std::env::var("POLARS_VISUALIZE_PHYSICAL_PLAN") {
            let visualization =
                crate::physical_plan::visualize_plan(root_phys_node, &phys_sm, expr_arena);
            std::fs::write(visual_path, visualization).unwrap();
        }

        let (mut graph, phys_to_graph) =
            crate::physical_plan::physical_plan_to_graph(root_phys_node, &phys_sm, expr_arena)?;

#[derive(Clone, Debug)]
pub enum PhysNodeKind {
    InMemorySource {
        df: Arc<DataFrame>,
        disable_morsel_split: bool,
    },

pub trait ComputeNode: Send {
    /// The name of this node.
    fn name(&self) -> &str;

    /// Update the state of this node given the state of our input and output
    /// ports. May be called multiple times until fully resolved for each
    /// execution phase.
    ///
    /// For each input pipe `recv` will contain a respective state of the
    /// send port that pipe is connected to when called, and it is expected when
    /// `update_state` returns it contains your computed receive port state.
    ///
    /// Similarly, for each output pipe `send` will contain the respective
    /// state of the input port that pipe is connected to when called, and you
    /// must update it to contain the desired state of your output port.
    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        state: &StreamingExecutionState,
    ) -> PolarsResult<()>;

/// 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>,
}

    let mut pipe_seq_offsets = SecondaryMap::new();
    loop {
        // Update the states.
        if polars_core::config::verbose() {
            eprintln!("polars-stream: updating graph state");
        }
        graph.update_all_states(&state, metrics.as_deref())?;

        // Find a subgraph to run.
        let (nodes, pipes) = find_runnable_subgraph(graph);

        if nodes.is_empty() {
            break;
        }

        // Run the subgraph until phase completion.
        run_subgraph(
            graph,
            &nodes,
            &pipes,
            &mut pipe_seq_offsets,
            &state,
            metrics.clone(),
        )?;
        ASYNC.block_in_place_on(async {
            // TODO: track this in metrics.
            while let Ok(handle) = subphase_tasks_recv.try_recv() {
                handle.await.unwrap()?;
            }
            PolarsResult::Ok(())
        })?;
        if polars_core::config::verbose() {
            eprintln!("polars-stream: done running graph phase");
        }
    }

/// Finds all runnable pipeline blockers in the graph, that is, nodes which:
///  - Only have blocked output ports.
///  - Have at least one ready input port connected to a ready output port.
fn find_runnable_pipeline_blockers(graph: &Graph) -> Vec<GraphNodeKey> {
    let mut blockers = Vec::new();
    for (node_key, node) in graph.nodes.iter() {
        // TODO: how does the multiplexer fit into this?
        let only_has_blocked_outputs = node
            .outputs
            .iter()
            .all(|o| graph.pipes[*o].send_state == PortState::Blocked);
        if !only_has_blocked_outputs {
            continue;
        }

        let has_input_ready = node.inputs.iter().any(|i| {
            graph.pipes[*i].send_state == PortState::Ready
                && graph.pipes[*i].recv_state == PortState::Ready
        });
        if has_input_ready {
            blockers.push(node_key);
        }
    }
    blockers
}

    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        state: &StreamingExecutionState,
    ) -> PolarsResult<()> {

        match &mut self.state {
            InMemoryJoinState::Sink { left, right, .. } => {
                left.update_state(&mut recv[0..1], &mut [], state)?;
                right.update_state(&mut recv[1..2], &mut [], state)?;
                send[0] = PortState::Blocked;
            },
            InMemoryJoinState::Source(source_node) => {
                recv[0] = PortState::Done;
                recv[1] = PortState::Done;
                source_node.update_state(&mut [], send, state)?;
            },
            InMemoryJoinState::Done => {
                recv[0] = PortState::Done;
                recv[1] = PortState::Done;
                send[0] = PortState::Done;
            },
        }
        Ok(())
    }

    fn is_memory_intensive_pipeline_blocker(&self) -> bool {
        matches!(self.state, InMemoryJoinState::Sink { .. })
    }

    /// Generic fallback for (as-of-yet) unsupported streaming mappings.
    /// Fully sinks all data to an in-memory data frame and uses the in-memory
    /// engine to perform the map.
    InMemoryMap {
        input: PhysStream,
        map: Arc<dyn DataFrameUdf>,
        /// A formatted string of what the in-memory map is. This usually calls format on the IR.
        format_str: Option<String>,
    },

            if options.maintain_order && options.keep_strategy == UniqueKeepStrategy::Last {
                // Unfortunately the order-preserving groupby always orders by the first occurrence
                // of the group so we can't lower this and have to fallback.
                let input_schema = phys_input.output_schema(phys_sm).clone();
                let lmdf = Arc::new(LateMaterializedDataFrame::default());

                        expr_arena,
                        phys_input.output_schema(phys_sm),

pub 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 { .. } | K::InMemoryJoin { .. } | K::ColumnarFunction { .. } => {
                Self::InMemoryFallback
            },
            K::InMemorySource { .. }
            | K::InputIndependentSelect { .. }
            | K::NegativeSlice { .. }
            | K::InMemorySink { .. }
            | K::Sort { .. }
            | K::GroupBy { .. }
            | K::EquiJoin { .. }
            | K::SemiAntiJoin { .. }
            | K::CrossJoin { .. }
            | K::Multiplexer { .. }
            | K::Gather { .. } => Self::MemoryIntensive,
            #[cfg(feature = "iejoin")]
            K::RangeJoin { .. } => Self::MemoryIntensive,
            #[cfg(feature = "merge_sorted")]
            K::MergeSorted { .. } => Self::MemoryIntensive,
            _ => Self::Generic,
        }
    }

        match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) {
            Ok(v) => Some(v),
            Err(e) => {
                // Fallback to normal engine if error is due to not being implemented
                // and auto_streaming is set, otherwise propagate error.
                if e.downcast_ref::<&str>()
                    .is_some_and(|s| s.starts_with("not yet implemented"))
                {
                    if polars_core::config::verbose() {
                        eprintln!(
                            "caught unimplemented error in new streaming engine, falling back to normal engine"
                        );
                    }
                    None
                } else {
                    std::panic::resume_unwind(e)
                }
            },
        }