polars-stream 0.55.2

use std::collections::VecDeque;

use polars_async::primitives::wait_group::WaitGroup;
use polars_ooc::{MostRecentSpillContext, ParameterFreeSpillContext, SpillFrame};
use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};

use super::compute_node_prelude::*;
use crate::morsel::SourceToken;

enum BufferedStream {
    Open(VecDeque<(SpillFrame, MorselSeq)>, MostRecentSpillContext),
    Closed,
}

impl BufferedStream {
    fn new() -> Self {
        Self::Open(
            VecDeque::new(),
            MostRecentSpillContext::new("multiplexer".into()),
        )
    }
}

pub struct MultiplexerNode {
    buffers: Vec<BufferedStream>,
}

impl MultiplexerNode {
    pub fn new() -> Self {
        Self {
            buffers: Vec::default(),
        }
    }
}

impl ComputeNode for MultiplexerNode {
    fn name(&self) -> &str {
        "multiplexer"
    }

    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        _state: &StreamingExecutionState,
    ) -> PolarsResult<()> {
        assert!(recv.len() == 1 && !send.is_empty());

        // Initialize buffered streams, and mark those for which the receiver
        // is no longer interested as closed.
        self.buffers.resize_with(send.len(), BufferedStream::new);
        for (s, b) in send.iter().zip(&mut self.buffers) {
            if *s == PortState::Done {
                *b = BufferedStream::Closed;
            }
        }

        // Check if either the input is done, or all outputs are done.
        let input_done = recv[0] == PortState::Done
            && self.buffers.iter().all(|b| match b {
                BufferedStream::Open(v, _) => v.is_empty(),
                BufferedStream::Closed => true,
            });
        let output_done = send.iter().all(|p| *p == PortState::Done);

        // If either side is done, everything is done.
        if input_done || output_done {
            recv[0] = PortState::Done;
            for s in send {
                *s = PortState::Done;
            }
            return Ok(());
        }

        let all_blocked = send.iter().all(|p| *p == PortState::Blocked);

        // Pass along the input state to the output.
        for (i, s) in send.iter_mut().enumerate() {
            let buffer_empty = match &self.buffers[i] {
                BufferedStream::Open(v, _) => v.is_empty(),
                BufferedStream::Closed => true,
            };
            *s = if buffer_empty && recv[0] == PortState::Done {
                PortState::Done
            } else if !buffer_empty || recv[0] == PortState::Ready {
                PortState::Ready
            } else {
                PortState::Blocked
            };
        }

        // We say we are ready to receive unless all outputs are blocked.
        recv[0] = if all_blocked {
            PortState::Blocked
        } else {
            PortState::Ready
        };
        Ok(())
    }

    fn spawn<'env, 's>(
        &'env mut self,
        scope: &'s TaskScope<'s, 'env>,
        recv_ports: &mut [Option<RecvPort<'_>>],
        send_ports: &mut [Option<SendPort<'_>>],
        _state: &'s StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert!(recv_ports.len() == 1 && !send_ports.is_empty());
        assert!(self.buffers.len() == send_ports.len());

        enum Listener<'a> {
            Active(
                UnboundedSender<(SpillFrame, MorselSeq, SourceToken)>,
                MostRecentSpillContext,
            ),
            Buffering(
                &'a mut VecDeque<(SpillFrame, MorselSeq)>,
                MostRecentSpillContext,
            ),
            Inactive,
        }

        let buffered_source_token = SourceToken::new();

        let (mut buf_senders, buf_receivers): (Vec<_>, Vec<_>) = self
            .buffers
            .iter_mut()
            .enumerate()
            .map(|(port_idx, buffer)| {
                if let BufferedStream::Open(buf, ctx) = buffer {
                    if send_ports[port_idx].is_some() {
                        let (rx, tx) = unbounded_channel();
                        (Listener::Active(rx, ctx.clone()), Some((buf, tx)))
                    } else {
                        (Listener::Buffering(buf, ctx.clone()), None)
                    }
                } else {
                    (Listener::Inactive, None)
                }
            })
            .unzip();
        // TODO: parallel multiplexing.
        if let Some(mut receiver) = recv_ports[0].take().map(|r| r.serial()) {
            let buffered_source_token = buffered_source_token.clone();
            join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                loop {
                    let Ok(mut morsel) = receiver.recv().await else {
                        break;
                    };
                    drop(morsel.take_consume_token());

                    let mut anyone_interested = false;
                    let mut active_listener_interested = false;
                    let seq = morsel.seq();
                    for buf_sender in &mut buf_senders {
                        match buf_sender {
                            Listener::Active(s, ctx) => {
                                let source_token = morsel.source_token().clone();
                                let sf = morsel.sf().clone();
                                ctx.register(&sf).await;
                                match s.send((sf, seq, source_token)) {
                                    Ok(_) => {
                                        anyone_interested = true;
                                        active_listener_interested = true;
                                    },
                                    Err(_) => *buf_sender = Listener::Inactive,
                                }
                            },
                            Listener::Buffering(b, ctx) => {
                                let sf = morsel.sf().clone();
                                ctx.register(&sf).await;
                                b.push_front((sf, seq));
                                anyone_interested = true;
                            },
                            Listener::Inactive => {},
                        }
                    }

                    if !anyone_interested {
                        break;
                    }

                    // If only buffering inputs are left, or we got a stop
                    // request from an input reading from old buffered data,
                    // request a stop from the source.
                    if !active_listener_interested || buffered_source_token.stop_requested() {
                        morsel.source_token().stop();
                    }
                }

                Ok(())
            }));
        }

        for (send_port, opt_buf_recv) in send_ports.iter_mut().zip(buf_receivers) {
            if let Some((buf, mut rx)) = opt_buf_recv {
                let mut sender = send_port.take().unwrap().serial();

                let wait_group = WaitGroup::default();
                let buffered_source_token = buffered_source_token.clone();
                join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                    // First we try to flush all the old buffered data.
                    while let Some((sf, seq)) = buf.pop_back() {
                        let mut morsel = Morsel::new(sf, seq, buffered_source_token.clone());
                        morsel.set_consume_token(wait_group.token());
                        if sender.send(morsel).await.is_err() {
                            return Ok(());
                        }

                        // Someone wants to stop, flush remainder into buffer.
                        if buffered_source_token.stop_requested() {
                            drop(sender);
                            while let Some((sf, seq, source_token)) = rx.recv().await {
                                source_token.stop();
                                buf.push_front((sf, seq));
                            }
                            return Ok(());
                        }

                        wait_group.wait().await;
                    }

                    // Then send along data from the multiplexer.
                    while let Some((sf, seq, source_token)) = rx.recv().await {
                        let mut morsel = Morsel::new(sf, seq, source_token);
                        morsel.set_consume_token(wait_group.token());
                        if sender.send(morsel).await.is_err() {
                            return Ok(());
                        }
                        wait_group.wait().await;
                    }
                    Ok(())
                }));
            }
        }
    }
}