laminar_core/streaming/sink/
mod.rs

1//! Sink — the single consumer of a source's bounded queue.
2
3use std::sync::Arc;
4
5use arrow::array::RecordBatch;
6use arrow::datatypes::SchemaRef;
7
8use super::channel::AsyncConsumer;
9use super::error::RecvError;
10use super::source::{Record, SourceMessage};
11
12/// The single consumer of a source's bounded queue.
13///
14/// Messages leave the queue only when the sink receives them, so a slow sink applies
15/// backpressure to producers (`ChannelFull`) instead of losing accepted input. Dropping the
16/// sink disconnects every producer.
17pub struct Sink<T: Record> {
18    consumer: AsyncConsumer<SourceMessage<T>>,
19    schema: SchemaRef,
20    closed: bool,
21}
22
23impl<T: Record> Sink<T> {
24    pub(crate) fn new(consumer: AsyncConsumer<SourceMessage<T>>, schema: SchemaRef) -> Self {
25        Self {
26            consumer,
27            schema,
28            closed: false,
29        }
30    }
31
32    /// Non-blocking poll. Returns the next batch.
33    /// Returns `None` on empty or closed queue. Check `is_disconnected()`
34    /// to distinguish.
35    pub fn poll(&mut self) -> Option<RecordBatch> {
36        match self.consumer.try_recv() {
37            Ok(message) => Some(message.into_batch()),
38            Err(crossfire::TryRecvError::Empty) => None,
39            Err(crossfire::TryRecvError::Disconnected) => {
40                self.closed = true;
41                None
42            }
43        }
44    }
45
46    /// Async receive. Awaits the next batch. Cancellation never loses a message.
47    ///
48    /// # Errors
49    ///
50    /// Returns `RecvError::Disconnected` once every producer has been dropped and the queue is
51    /// empty.
52    pub async fn recv_async(&mut self) -> Result<RecordBatch, RecvError> {
53        let Ok(message) = self.consumer.recv().await else {
54            self.closed = true;
55            return Err(RecvError::Disconnected);
56        };
57        Ok(message.into_batch())
58    }
59
60    /// Returns true if the queue has been observed closed and drained.
61    #[must_use]
62    pub fn is_disconnected(&self) -> bool {
63        self.closed
64    }
65
66    /// Returns the schema for this sink.
67    #[must_use]
68    pub fn schema(&self) -> SchemaRef {
69        Arc::clone(&self.schema)
70    }
71}
72
73impl<T: Record + std::fmt::Debug> std::fmt::Debug for Sink<T> {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("Sink")
76            .field("closed", &self.closed)
77            .finish_non_exhaustive()
78    }
79}
80
81#[cfg(test)]
82mod tests;
laminar_core/streaming/sink/mod.rs

laminar_core/streaming/sink/
mod.rs
