micromegas_analytics/lakehouse/
table_scan_rewrite.rs

1use crate::{lakehouse::materialized_view::MaterializedView, time::TimeRange};
2use datafusion::error::DataFusionError;
3use datafusion::logical_expr::Filter;
4use datafusion::logical_expr::utils::conjunction;
5use datafusion::{
6    common::tree_node::Transformed, config::ConfigOptions, datasource::DefaultTableSource,
7    logical_expr::LogicalPlan, optimizer::AnalyzerRule,
8};
9use std::sync::Arc;
10
11/// An analyzer rule that rewrites table scans to include time-based filters.
12#[derive(Debug)]
13pub struct TableScanRewrite {
14    query_range: TimeRange,
15}
16
17impl TableScanRewrite {
18    pub fn new(query_range: TimeRange) -> Self {
19        Self { query_range }
20    }
21
22    fn rewrite_plan(
23        &self,
24        plan: LogicalPlan,
25        _options: &ConfigOptions,
26    ) -> datafusion::error::Result<Transformed<LogicalPlan>> {
27        if let LogicalPlan::TableScan(ts) = &plan {
28            let table_source = ts
29                .source
30                .downcast_ref::<DefaultTableSource>()
31                .ok_or_else(|| {
32                    DataFusionError::Execution(String::from(
33                        "error casting table source as DefaultTableSource",
34                    ))
35                })?;
36            // Only rewrite MaterializedView tables, skip others (like table functions)
37            let Some(mat_view) = table_source
38                .table_provider
39                .downcast_ref::<MaterializedView>()
40            else {
41                // This is not a MaterializedView (e.g., a table function), skip rewriting
42                return Ok(Transformed::no(plan));
43            };
44            let view = mat_view.get_view();
45            let filters = view
46                .make_time_filter(self.query_range.begin, self.query_range.end)
47                .map_err(|e| DataFusionError::External(e.into()))?;
48            let pred = conjunction(filters).ok_or_else(|| {
49                DataFusionError::Execution(String::from("error making a conjunction"))
50            })?;
51            let filter = Filter::try_new(pred, Arc::new(plan.clone()))?;
52            Ok(Transformed::yes(LogicalPlan::Filter(filter)))
53        } else {
54            Ok(Transformed::no(plan))
55        }
56    }
57}
58
59impl AnalyzerRule for TableScanRewrite {
60    fn name(&self) -> &str {
61        "table_scan_rewrite"
62    }
63
64    fn analyze(
65        &self,
66        plan: LogicalPlan,
67        options: &ConfigOptions,
68    ) -> datafusion::error::Result<LogicalPlan> {
69        plan.transform_up_with_subqueries(|plan| self.rewrite_plan(plan, options))
70            .map(|res| res.data)
71    }
72}
micromegas_analytics/lakehouse/table_scan_rewrite.rs

micromegas_analytics/lakehouse/
table_scan_rewrite.rs
