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::prelude::Expr;
10use std::sync::Arc;
11
12/// A DataFusion `TableProvider` for a log stream.
13#[derive(Debug)]
14pub struct LogStreamTableProvider {
15    /// The underlying log stream execution plan.
16    pub log_stream: Arc<TaskLogExecPlan>,
17}
18
19#[async_trait]
20impl TableProvider for LogStreamTableProvider {
21    fn schema(&self) -> SchemaRef {
22        self.log_stream.schema()
23    }
24
25    fn table_type(&self) -> TableType {
26        TableType::Temporary
27    }
28
29    async fn scan(
30        &self,
31        _state: &dyn Session,
32        _projection: Option<&Vec<usize>>,
33        _filters: &[Expr],
34        limit: Option<usize>,
35    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
36        // Wrap the execution plan in a GlobalLimitExec if a limit is provided.
37        // DataFusion trusts us to apply the limit - if we ignore it, too many rows
38        // will be returned to the client.
39        let plan: Arc<dyn ExecutionPlan> = self.log_stream.clone();
40        if let Some(fetch) = limit {
41            Ok(Arc::new(GlobalLimitExec::new(plan, 0, Some(fetch))))
42        } else {
43            Ok(plan)
44        }
45    }
46}
micromegas_analytics/dfext/log_stream_table_provider.rs

micromegas_analytics/dfext/
log_stream_table_provider.rs
