micromegas_analytics/dfext/
async_log_stream.rs

1use datafusion::{
2    arrow::{
3        array::{PrimitiveBuilder, RecordBatch, StringBuilder},
4        datatypes::{SchemaRef, TimestampNanosecondType},
5    },
6    common::Result,
7    error::DataFusionError,
8    execution::RecordBatchStream,
9};
10use futures::Stream;
11use std::{
12    sync::Arc,
13    task::{Context, Poll},
14};
15use tokio::sync::mpsc;
16
17/// A stream of log messages that can be converted into a `RecordBatchStream`.
18///
19/// The channel carries `Result<(time, msg), String>`: an `Err` item ends the stream with a
20/// genuine `RecordBatchStream` error (propagating through `execute_stream`/`collect` as a query
21/// execution failure) instead of being folded into one more `(time, msg)` log row.
22pub struct AsyncLogStream {
23    schema: SchemaRef,
24    rx: mpsc::Receiver<Result<(chrono::DateTime<chrono::Utc>, String), String>>,
25}
26
27impl AsyncLogStream {
28    pub fn new(
29        schema: SchemaRef,
30        rx: mpsc::Receiver<Result<(chrono::DateTime<chrono::Utc>, String), String>>,
31    ) -> Self {
32        Self { schema, rx }
33    }
34}
35
36impl Stream for AsyncLogStream {
37    type Item = Result<RecordBatch>;
38
39    fn poll_next(
40        mut self: std::pin::Pin<&mut Self>,
41        cx: &mut Context<'_>,
42    ) -> Poll<Option<Self::Item>> {
43        let mut messages = vec![];
44        let limit = self.rx.max_capacity();
45        if self
46            .rx
47            .poll_recv_many(cx, &mut messages, limit)
48            .is_pending()
49        {
50            cx.waker().wake_by_ref();
51            return Poll::Pending;
52        }
53        if messages.is_empty() {
54            if self.rx.is_closed() {
55                // channel closed, aborting
56                return Poll::Ready(None);
57            }
58            // not sure this can happen
59            cx.waker().wake_by_ref();
60            return Poll::Pending;
61        }
62
63        // An Err item ends the stream as a query error. Ok items received in the same poll batch
64        // ahead of the Err are dropped -- acceptable, they are transient progress lines on a
65        // query that is failing anyway.
66        let mut times = PrimitiveBuilder::<TimestampNanosecondType>::with_capacity(messages.len());
67        let mut msgs = StringBuilder::new();
68        for msg in messages {
69            match msg {
70                Ok((time, text)) => {
71                    times.append_value(time.timestamp_nanos_opt().unwrap_or_default());
72                    msgs.append_value(text);
73                }
74                Err(err_msg) => {
75                    return Poll::Ready(Some(Err(DataFusionError::Internal(err_msg))));
76                }
77            }
78        }
79
80        let rb_res = RecordBatch::try_new(
81            self.schema.clone(),
82            vec![
83                Arc::new(times.finish().with_timezone_utc()),
84                Arc::new(msgs.finish()),
85            ],
86        )
87        .map_err(|e| DataFusionError::ArrowError(e.into(), None));
88        Poll::Ready(Some(rb_res))
89    }
90
91    fn size_hint(&self) -> (usize, Option<usize>) {
92        (self.rx.len(), Some(self.rx.len()))
93    }
94}
95
96impl RecordBatchStream for AsyncLogStream {
97    fn schema(&self) -> SchemaRef {
98        Arc::clone(&self.schema)
99    }
100}
micromegas_analytics/dfext/async_log_stream.rs

micromegas_analytics/dfext/
async_log_stream.rs
