micromegas_analytics/dfext/
task_log_exec_plan.rs

1use datafusion::arrow::datatypes::DataType;
2use datafusion::arrow::datatypes::Field;
3use datafusion::arrow::datatypes::Schema;
4use datafusion::arrow::datatypes::SchemaRef;
5use datafusion::arrow::datatypes::TimeUnit;
6use datafusion::common::Statistics;
7use datafusion::common::internal_err;
8use datafusion::error::DataFusionError;
9use datafusion::execution::SendableRecordBatchStream;
10use datafusion::execution::TaskContext;
11use datafusion::physical_expr::EquivalenceProperties;
12use datafusion::physical_plan::DisplayAs;
13use datafusion::physical_plan::DisplayFormatType;
14use datafusion::physical_plan::ExecutionPlan;
15use datafusion::physical_plan::Partitioning;
16use datafusion::physical_plan::PlanProperties;
17use datafusion::physical_plan::execution_plan::Boundedness;
18use datafusion::physical_plan::execution_plan::EmissionType;
19use std::sync::Arc;
20use tokio::sync::mpsc;
21
22use super::async_log_stream::AsyncLogStream;
23
24/// A type alias for a function that spawns a log message receiver.
25///
26/// The channel item is `Result<(time, msg), String>` so a query-level failure can be surfaced as
27/// a genuine `RecordBatchStream` error (see `AsyncLogStream::poll_next`) instead of being
28/// swallowed as one more log row. Existing producers only ever send `Ok` items, via
29/// `LogSender::write_log_entry`.
30pub type TaskSpawner = dyn FnOnce() -> mpsc::Receiver<Result<(chrono::DateTime<chrono::Utc>, String), String>>
31    + Sync
32    + Send;
33
34/// An `ExecutionPlan` that provides a stream of log messages.
35pub struct TaskLogExecPlan {
36    schema: SchemaRef,
37    cache: Arc<PlanProperties>,
38    spawner: std::sync::Mutex<Option<Box<TaskSpawner>>>,
39}
40
41impl DisplayAs for TaskLogExecPlan {
42    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
43        match t {
44            DisplayFormatType::Default
45            | DisplayFormatType::Verbose
46            | DisplayFormatType::TreeRender => {
47                write!(f, "TaskLogExecPlan")
48            }
49        }
50    }
51}
52
53impl std::fmt::Debug for TaskLogExecPlan {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        write!(f, "TaskLogExecPlan")
56    }
57}
58
59impl TaskLogExecPlan {
60    pub fn new(spawner: Box<TaskSpawner>) -> Self {
61        let schema = Arc::new(Schema::new(vec![
62            Field::new(
63                "time",
64                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
65                false,
66            ),
67            Field::new("msg", DataType::Utf8, false),
68        ]));
69
70        let cache = PlanProperties::new(
71            EquivalenceProperties::new(Arc::clone(&schema)),
72            Partitioning::RoundRobinBatch(1),
73            EmissionType::Incremental,
74            Boundedness::Unbounded {
75                requires_infinite_memory: false,
76            },
77        );
78
79        Self {
80            schema,
81            cache: Arc::new(cache),
82            spawner: std::sync::Mutex::new(Some(spawner)),
83        }
84    }
85}
86
87impl ExecutionPlan for TaskLogExecPlan {
88    fn name(&self) -> &'static str {
89        "LogExecPlan"
90    }
91
92    fn schema(&self) -> SchemaRef {
93        self.schema.clone()
94    }
95
96    fn properties(&self) -> &Arc<PlanProperties> {
97        &self.cache
98    }
99
100    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
101        vec![]
102    }
103
104    fn with_new_children(
105        self: Arc<Self>,
106        children: Vec<Arc<dyn ExecutionPlan>>,
107    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
108        if children.is_empty() {
109            Ok(self)
110        } else {
111            internal_err!("Children cannot be replaced in LogExecPlan")
112        }
113    }
114
115    fn execute(
116        &self,
117        partition: usize,
118        _context: Arc<TaskContext>,
119    ) -> datafusion::error::Result<SendableRecordBatchStream> {
120        if partition >= 1 {
121            return internal_err!("Invalid partition {partition} for LogExecPlan");
122        }
123
124        let mut spawner = self.spawner.lock().map_err(|_| {
125            DataFusionError::Execution("Error locking mutex in LogExecPlan".to_owned())
126        })?;
127        if let Some(fun) = spawner.take() {
128            drop(spawner);
129            Ok(Box::pin(AsyncLogStream::new(self.schema.clone(), fun())))
130        } else {
131            internal_err!("Spawner already taken in LogExecPlan")
132        }
133    }
134
135    fn partition_statistics(
136        &self,
137        _partition: Option<usize>,
138    ) -> datafusion::error::Result<Arc<Statistics>> {
139        Ok(Arc::new(Statistics::new_unknown(&self.schema)))
140    }
141}
micromegas_analytics/dfext/task_log_exec_plan.rs

micromegas_analytics/dfext/
task_log_exec_plan.rs
