micromegas_analytics/dfext/
log_stream_table_provider.rs

1use super::task_log_exec_plan::TaskLogExecPlan;
2use async_trait::async_trait;
3use datafusion::arrow::datatypes::SchemaRef;
4use datafusion::catalog::Session;
5use datafusion::catalog::TableProvider;
6use datafusion::datasource::TableType;
7use datafusion::physical_plan::ExecutionPlan;
8use datafusion::physical_plan::limit::GlobalLimitExec;
9use datafusion::physical_plan::projection::ProjectionExec;
10use datafusion::prelude::Expr;
11use std::sync::Arc;
12
13/// A DataFusion `TableProvider` for a log stream.
14#[derive(Debug)]
15pub struct LogStreamTableProvider {
16    /// The underlying log stream execution plan.
17    pub log_stream: Arc<TaskLogExecPlan>,
18}
19
20#[async_trait]
21impl TableProvider for LogStreamTableProvider {
22    fn schema(&self) -> SchemaRef {
23        self.log_stream.schema()
24    }
25
26    fn table_type(&self) -> TableType {
27        TableType::Temporary
28    }
29
30    async fn scan(
31        &self,
32        _state: &dyn Session,
33        projection: Option<&Vec<usize>>,
34        _filters: &[Expr],
35        limit: Option<usize>,
36    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
37        let mut plan: Arc<dyn ExecutionPlan> = self.log_stream.clone();
38        // DataFusion trusts the scan to apply the pushed-down projection: returning the full
39        // (time, msg) schema against a narrower projection fails physical planning (e.g.
40        // `SELECT msg FROM retire_partitions(...)` or `SELECT msg AS rule_id FROM
41        // deny_queries(...)`). Same idiom as `process_spans_table_function.rs`.
42        if let Some(projection) = projection {
43            let schema = plan.schema();
44            let projected_exprs: Vec<(Arc<dyn datafusion::physical_expr::PhysicalExpr>, String)> =
45                projection
46                    .iter()
47                    .map(|&i| {
48                        let name = schema.field(i).name().clone();
49                        let expr = Arc::new(datafusion::physical_expr::expressions::Column::new(
50                            &name, i,
51                        ))
52                            as Arc<dyn datafusion::physical_expr::PhysicalExpr>;
53                        (expr, name)
54                    })
55                    .collect();
56            plan = Arc::new(ProjectionExec::try_new(projected_exprs, plan)?);
57        }
58        // DataFusion likewise trusts us to apply the limit - if we ignore it, too many rows
59        // will be returned to the client.
60        if let Some(fetch) = limit {
61            plan = Arc::new(GlobalLimitExec::new(plan, 0, Some(fetch)));
62        }
63        Ok(plan)
64    }
65}
micromegas_analytics/dfext/log_stream_table_provider.rs

micromegas_analytics/dfext/
log_stream_table_provider.rs
