laminar_db/
query_rows.rs

1//! The rows of an ad-hoc query, delivered in order through a bounded queue. The query runs only as
2//! fast as its reader reads, and a result that ends early ends with an error, never quietly.
3
4use arrow::array::RecordBatch;
5use arrow::datatypes::SchemaRef;
6use datafusion::physical_plan::SendableRecordBatchStream;
7use tokio::sync::mpsc;
8use tokio_stream::StreamExt;
9use tokio_util::sync::CancellationToken;
10
11use crate::error::DbError;
12
13/// Batches a result holds ahead of its reader.
14const QUEUED_BATCHES: usize = 8;
15
16enum Delivery {
17    Batch(RecordBatch),
18    Failed(DbError),
19    End,
20}
21
22impl std::fmt::Debug for Delivery {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        match self {
25            Self::Batch(batch) => write!(f, "Batch({} rows)", batch.num_rows()),
26            Self::Failed(error) => write!(f, "Failed({error})"),
27            Self::End => f.write_str("End"),
28        }
29    }
30}
31
32/// The rows of an ad-hoc query, in order. Reading them to the end returns every row the query
33/// produced; a query that fails, or is cancelled before its last row, ends them with an error.
34#[derive(Debug)]
35pub struct QueryRows {
36    schema: SchemaRef,
37    receiver: mpsc::Receiver<Delivery>,
38    ended: bool,
39}
40
41impl QueryRows {
42    /// Run `stream` on a task that queues its batches for the returned rows. The task stops at the
43    /// stream's end or first error, on `cancel`, or when the rows are dropped, and then calls
44    /// `finished`.
45    pub(crate) fn spawn(
46        mut stream: SendableRecordBatchStream,
47        cancel: CancellationToken,
48        finished: impl FnOnce() + Send + 'static,
49    ) -> Self {
50        let schema = stream.schema();
51        let (sender, receiver) = mpsc::channel(QUEUED_BATCHES);
52        tokio::spawn(async move {
53            loop {
54                let delivery = tokio::select! {
55                    () = cancel.cancelled() => break,
56                    next = stream.next() => match next {
57                        Some(Ok(batch)) => Delivery::Batch(batch),
58                        Some(Err(error)) => Delivery::Failed(execution_error(error)),
59                        None => Delivery::End,
60                    },
61                };
62                let last = !matches!(delivery, Delivery::Batch(_));
63                let sent = tokio::select! {
64                    () = cancel.cancelled() => break,
65                    sent = sender.send(delivery) => sent,
66                };
67                if last || sent.is_err() {
68                    break;
69                }
70            }
71            finished();
72        });
73        Self {
74            schema,
75            receiver,
76            ended: false,
77        }
78    }
79
80    /// Rows that were already read to the end.
81    #[cfg(feature = "api")]
82    pub(crate) fn ready(schema: SchemaRef, batches: Vec<RecordBatch>) -> Self {
83        let (sender, receiver) = mpsc::channel(batches.len() + 1);
84        for delivery in batches
85            .into_iter()
86            .map(Delivery::Batch)
87            .chain(std::iter::once(Delivery::End))
88        {
89            // The channel holds every delivery and its receiver is alive.
90            let _ = sender.try_send(delivery);
91        }
92        Self {
93            schema,
94            receiver,
95            ended: false,
96        }
97    }
98
99    /// Schema of every batch.
100    #[must_use]
101    pub fn schema(&self) -> SchemaRef {
102        SchemaRef::clone(&self.schema)
103    }
104
105    /// The next batch, waiting for it; `None` after the last one.
106    ///
107    /// # Errors
108    /// The query's error if it failed, or [`DbError::InvalidOperation`] if it was cancelled
109    /// before its last row.
110    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>, DbError> {
111        if self.ended {
112            return Ok(None);
113        }
114        let delivery = self.receiver.recv().await;
115        self.deliver(delivery)
116    }
117
118    /// The next batch if one is ready; `None` when none is ready yet or after the last one.
119    ///
120    /// # Errors
121    /// As [`Self::next_batch`].
122    pub fn try_next_batch(&mut self) -> Result<Option<RecordBatch>, DbError> {
123        if self.ended {
124            return Ok(None);
125        }
126        match self.receiver.try_recv() {
127            Ok(delivery) => self.deliver(Some(delivery)),
128            Err(mpsc::error::TryRecvError::Empty) => Ok(None),
129            Err(mpsc::error::TryRecvError::Disconnected) => self.deliver(None),
130        }
131    }
132
133    /// The next batch, blocking the thread until it is ready; for callers outside an async
134    /// runtime.
135    ///
136    /// # Errors
137    /// As [`Self::next_batch`].
138    ///
139    /// # Panics
140    /// When called from within an async runtime.
141    pub fn blocking_next_batch(&mut self) -> Result<Option<RecordBatch>, DbError> {
142        if self.ended {
143            return Ok(None);
144        }
145        let delivery = self.receiver.blocking_recv();
146        self.deliver(delivery)
147    }
148
149    /// Every remaining batch.
150    ///
151    /// # Errors
152    /// As [`Self::next_batch`].
153    pub async fn read_all(mut self) -> Result<Vec<RecordBatch>, DbError> {
154        let mut batches = Vec::new();
155        while let Some(batch) = self.next_batch().await? {
156            batches.push(batch);
157        }
158        Ok(batches)
159    }
160
161    fn deliver(&mut self, delivery: Option<Delivery>) -> Result<Option<RecordBatch>, DbError> {
162        match delivery {
163            Some(Delivery::Batch(batch)) => Ok(Some(batch)),
164            Some(Delivery::End) => {
165                self.ended = true;
166                Ok(None)
167            }
168            Some(Delivery::Failed(error)) => {
169                self.ended = true;
170                Err(error)
171            }
172            None => {
173                self.ended = true;
174                Err(DbError::InvalidOperation(
175                    "the query was cancelled before its last row".into(),
176                ))
177            }
178        }
179    }
180}
181
182/// A query fails with the [`DbError`] a table scan raised, which `DataFusion` carries as an
183/// external error, or with `DataFusion`'s own error.
184pub(crate) fn execution_error(error: datafusion_common::DataFusionError) -> DbError {
185    match error {
186        datafusion_common::DataFusionError::External(inner) => match inner.downcast::<DbError>() {
187            Ok(error) => *error,
188            Err(inner) => datafusion_common::DataFusionError::External(inner).into(),
189        },
190        error => error.into(),
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use std::sync::Arc;
197
198    use arrow::array::Int64Array;
199    use arrow::datatypes::{DataType, Field, Schema};
200    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
201
202    use super::*;
203
204    fn batches(count: i64) -> (SchemaRef, Vec<RecordBatch>) {
205        let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int64, false)]));
206        let batches = (0..count)
207            .map(|n| {
208                RecordBatch::try_new(
209                    Arc::clone(&schema),
210                    vec![Arc::new(Int64Array::from(vec![n]))],
211                )
212                .unwrap()
213            })
214            .collect();
215        (schema, batches)
216    }
217
218    fn rows_of(
219        schema: SchemaRef,
220        items: Vec<Result<RecordBatch, datafusion_common::DataFusionError>>,
221        cancel: CancellationToken,
222    ) -> QueryRows {
223        let stream = RecordBatchStreamAdapter::new(schema, futures::stream::iter(items));
224        QueryRows::spawn(Box::pin(stream), cancel, || {})
225    }
226
227    /// A reader that falls far behind still receives every batch: the query waits for it.
228    #[tokio::test]
229    async fn a_slow_reader_receives_every_batch() {
230        let (schema, batches) = batches(5_000);
231        let rows = rows_of(
232            schema,
233            batches.into_iter().map(Ok).collect(),
234            CancellationToken::new(),
235        );
236        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
237        let read = rows.read_all().await.unwrap();
238        let values: Vec<i64> = read
239            .iter()
240            .map(|batch| {
241                batch
242                    .column(0)
243                    .as_any()
244                    .downcast_ref::<Int64Array>()
245                    .unwrap()
246                    .value(0)
247            })
248            .collect();
249        assert_eq!(values, (0..5_000).collect::<Vec<_>>());
250    }
251
252    #[tokio::test]
253    async fn a_failure_or_a_cancellation_ends_the_rows_with_an_error() {
254        let (schema, mut batches) = batches(2);
255        let mut failed = rows_of(
256            Arc::clone(&schema),
257            vec![
258                Ok(batches.remove(0)),
259                Err(datafusion_common::DataFusionError::External(Box::new(
260                    DbError::MaterializedView("scan failed".into()),
261                ))),
262                Ok(batches.remove(0)),
263            ],
264            CancellationToken::new(),
265        );
266        assert!(failed.next_batch().await.unwrap().is_some());
267        let error = failed.next_batch().await.unwrap_err();
268        assert!(matches!(error, DbError::MaterializedView(_)), "{error}");
269        assert!(failed.next_batch().await.unwrap().is_none());
270
271        let cancel = CancellationToken::new();
272        let pending = RecordBatchStreamAdapter::new(schema, futures::stream::pending());
273        let mut cancelled = QueryRows::spawn(Box::pin(pending), cancel.clone(), || {});
274        cancel.cancel();
275        let error = cancelled.next_batch().await.unwrap_err();
276        assert!(error.to_string().contains("cancelled"), "{error}");
277    }
278}
laminar_db/query_rows.rs

laminar_db/
query_rows.rs
