micromegas_analytics/lakehouse/
partitioned_table_provider.rs

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

micromegas_analytics/lakehouse/
partitioned_table_provider.rs
