micromegas_analytics/lakehouse/
partitioned_table_provider.rs

1use super::{
2    partition::Partition,
3    partitioned_execution_plan::{ScanOrdering, make_partitioned_execution_plan},
4    reader_factory::ReaderFactory,
5};
6use async_trait::async_trait;
7use datafusion::{
8    arrow::datatypes::SchemaRef,
9    catalog::{Session, TableProvider},
10    datasource::TableType,
11    logical_expr::TableProviderFilterPushDown,
12    physical_plan::ExecutionPlan,
13    prelude::*,
14};
15use std::sync::Arc;
16
17/// A DataFusion `TableProvider` for a set of pre-defined partitions.
18pub struct PartitionedTableProvider {
19    schema: SchemaRef,
20    reader_factory: Arc<ReaderFactory>,
21    partitions: Arc<Vec<Partition>>,
22    scan_ordering: ScanOrdering,
23}
24
25impl std::fmt::Debug for PartitionedTableProvider {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        f.debug_struct("PartitionedTableProvider")
28            .field("schema", &self.schema)
29            .field("partitions_count", &self.partitions.len())
30            .finish()
31    }
32}
33
34impl PartitionedTableProvider {
35    pub fn new(
36        schema: SchemaRef,
37        reader_factory: Arc<ReaderFactory>,
38        partitions: Arc<Vec<Partition>>,
39    ) -> Self {
40        Self {
41            schema,
42            reader_factory,
43            partitions,
44            scan_ordering: ScanOrdering::Unordered,
45        }
46    }
47
48    /// Builds a `PartitionedTableProvider` that declares `scan_ordering` as an ordering the
49    /// scan's rows already satisfy (see `make_partitioned_execution_plan`).
50    pub fn with_scan_ordering(
51        schema: SchemaRef,
52        reader_factory: Arc<ReaderFactory>,
53        partitions: Arc<Vec<Partition>>,
54        scan_ordering: ScanOrdering,
55    ) -> Self {
56        Self {
57            schema,
58            reader_factory,
59            partitions,
60            scan_ordering,
61        }
62    }
63}
64
65#[async_trait]
66impl TableProvider for PartitionedTableProvider {
67    fn schema(&self) -> SchemaRef {
68        self.schema.clone()
69    }
70
71    fn table_type(&self) -> TableType {
72        TableType::Base
73    }
74
75    async fn scan(
76        &self,
77        state: &dyn Session,
78        projection: Option<&Vec<usize>>,
79        filters: &[Expr],
80        limit: Option<usize>,
81    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
82        make_partitioned_execution_plan(
83            self.schema(),
84            self.reader_factory.clone(),
85            state,
86            projection,
87            filters,
88            limit,
89            self.partitions.clone(),
90            &self.scan_ordering,
91        )
92    }
93
94    /// Tell DataFusion to push filters down to the scan method
95    fn supports_filters_pushdown(
96        &self,
97        filters: &[&Expr],
98    ) -> datafusion::error::Result<Vec<TableProviderFilterPushDown>> {
99        // Inexact because the pruning can't handle all expressions and pruning
100        // is not done at the row level -- there may be rows in returned files
101        // that do not pass the filter
102        Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
103    }
104}
micromegas_analytics/lakehouse/partitioned_table_provider.rs

micromegas_analytics/lakehouse/
partitioned_table_provider.rs
