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 -- see
22/// `tasks/blocks_view_ordered_merges_plan.md`'s Design §3.
23pub struct AsyncLogStream {
24    schema: SchemaRef,
25    rx: mpsc::Receiver<Result<(chrono::DateTime<chrono::Utc>, String), String>>,
26}
27
28impl AsyncLogStream {
29    pub fn new(
30        schema: SchemaRef,
31        rx: mpsc::Receiver<Result<(chrono::DateTime<chrono::Utc>, String), String>>,
32    ) -> Self {
33        Self { schema, rx }
34    }
35}
36
37impl Stream for AsyncLogStream {
38    type Item = Result<RecordBatch>;
39
40    fn poll_next(
41        mut self: std::pin::Pin<&mut Self>,
42        cx: &mut Context<'_>,
43    ) -> Poll<Option<Self::Item>> {
44        let mut messages = vec![];
45        let limit = self.rx.max_capacity();
46        if self
47            .rx
48            .poll_recv_many(cx, &mut messages, limit)
49            .is_pending()
50        {
51            cx.waker().wake_by_ref();
52            return Poll::Pending;
53        }
54        if messages.is_empty() {
55            if self.rx.is_closed() {
56                // channel closed, aborting
57                return Poll::Ready(None);
58            }
59            // not sure this can happen
60            cx.waker().wake_by_ref();
61            return Poll::Pending;
62        }
63
64        // An Err item ends the stream as a query error. Ok items received in the same poll batch
65        // ahead of the Err are dropped -- acceptable, they are transient progress lines on a
66        // query that is failing anyway.
67        let mut times = PrimitiveBuilder::<TimestampNanosecondType>::with_capacity(messages.len());
68        let mut msgs = StringBuilder::new();
69        for msg in messages {
70            match msg {
71                Ok((time, text)) => {
72                    times.append_value(time.timestamp_nanos_opt().unwrap_or_default());
73                    msgs.append_value(text);
74                }
75                Err(err_msg) => {
76                    return Poll::Ready(Some(Err(DataFusionError::Execution(err_msg))));
77                }
78            }
79        }
80
81        let rb_res = RecordBatch::try_new(
82            self.schema.clone(),
83            vec![
84                Arc::new(times.finish().with_timezone_utc()),
85                Arc::new(msgs.finish()),
86            ],
87        )
88        .map_err(|e| DataFusionError::ArrowError(e.into(), None));
89        Poll::Ready(Some(rb_res))
90    }
91
92    fn size_hint(&self) -> (usize, Option<usize>) {
93        (self.rx.len(), Some(self.rx.len()))
94    }
95}
96
97impl RecordBatchStream for AsyncLogStream {
98    fn schema(&self) -> SchemaRef {
99        Arc::clone(&self.schema)
100    }
101}
micromegas_analytics/dfext/async_log_stream.rs

micromegas_analytics/dfext/
async_log_stream.rs
