micromegas_analytics/lakehouse/
list_partitions_table_function.rs

1use crate::sql_arrow_bridge::rows_to_record_batch;
2use async_trait::async_trait;
3use datafusion::arrow::datatypes::DataType;
4use datafusion::arrow::datatypes::Field;
5use datafusion::arrow::datatypes::Schema;
6use datafusion::arrow::datatypes::SchemaRef;
7use datafusion::arrow::datatypes::TimeUnit;
8use datafusion::catalog::Session;
9use datafusion::catalog::TableFunctionArgs;
10use datafusion::catalog::TableFunctionImpl;
11use datafusion::catalog::TableProvider;
12use datafusion::datasource::TableType;
13use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
14use datafusion::error::DataFusionError;
15use datafusion::physical_plan::ExecutionPlan;
16use datafusion::prelude::Expr;
17use micromegas_ingestion::data_lake_connection::DataLakeConnection;
18use micromegas_tracing::prelude::*;
19use std::sync::Arc;
20
21/// A DataFusion `TableFunctionImpl` for listing lakehouse partitions.
22#[derive(Debug)]
23pub struct ListPartitionsTableFunction {
24    lake: Arc<DataLakeConnection>,
25}
26
27impl ListPartitionsTableFunction {
28    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
29        Self { lake }
30    }
31}
32
33impl TableFunctionImpl for ListPartitionsTableFunction {
34    fn call_with_args(
35        &self,
36        _args: TableFunctionArgs,
37    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
38        Ok(Arc::new(ListPartitionsTableProvider {
39            lake: self.lake.clone(),
40        }))
41    }
42}
43
44/// A DataFusion `TableProvider` for listing lakehouse partitions.
45#[derive(Debug)]
46pub struct ListPartitionsTableProvider {
47    pub lake: Arc<DataLakeConnection>,
48}
49
50#[async_trait]
51impl TableProvider for ListPartitionsTableProvider {
52    fn schema(&self) -> SchemaRef {
53        Arc::new(Schema::new(vec![
54            Field::new("view_set_name", DataType::Utf8, false),
55            Field::new("view_instance_id", DataType::Utf8, false),
56            Field::new(
57                "begin_insert_time",
58                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
59                false,
60            ),
61            Field::new(
62                "end_insert_time",
63                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
64                false,
65            ),
66            Field::new(
67                "min_event_time",
68                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
69                true,
70            ),
71            Field::new(
72                "max_event_time",
73                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
74                true,
75            ),
76            Field::new(
77                "updated",
78                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
79                false,
80            ),
81            Field::new("file_path", DataType::Utf8, true),
82            Field::new("file_size", DataType::Int64, false),
83            Field::new("file_schema_hash", DataType::Binary, false),
84            Field::new("source_data_hash", DataType::Binary, false),
85            Field::new("num_rows", DataType::Int64, false),
86            Field::new("partition_format_version", DataType::Int32, false),
87            Field::new(
88                "sort_order",
89                DataType::List(Arc::new(Field::new("tag", DataType::Utf8, false))),
90                true,
91            ),
92        ]))
93    }
94
95    fn table_type(&self) -> TableType {
96        TableType::Temporary
97    }
98
99    async fn scan(
100        &self,
101        _state: &dyn Session,
102        projection: Option<&Vec<usize>>,
103        _filters: &[Expr],
104        limit: Option<usize>,
105    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
106        // Build query with optional LIMIT clause pushed down to PostgreSQL.
107        // DataFusion only pushes the limit when it's safe to do so (i.e., when there
108        // are no WHERE clauses that could filter rows). When filters are present,
109        // DataFusion passes limit=None and applies the limit after filtering.
110        // Important: DataFusion trusts us to apply the limit - if we ignore it,
111        // too many rows will be returned to the client.
112        let query = if let Some(n) = limit {
113            format!(
114                "SELECT view_set_name,
115                        view_instance_id,
116                        begin_insert_time,
117                        end_insert_time,
118                        min_event_time,
119                        max_event_time,
120                        updated,
121                        file_path,
122                        file_size,
123                        file_schema_hash,
124                        source_data_hash,
125                        num_rows,
126                        partition_format_version,
127                        sort_order
128                 FROM lakehouse_partitions
129                 LIMIT {n};"
130            )
131        } else {
132            "SELECT view_set_name,
133                    view_instance_id,
134                    begin_insert_time,
135                    end_insert_time,
136                    min_event_time,
137                    max_event_time,
138                    updated,
139                    file_path,
140                    file_size,
141                    file_schema_hash,
142                    source_data_hash,
143                    num_rows,
144                    partition_format_version,
145                    sort_order
146             FROM lakehouse_partitions;"
147                .to_string()
148        };
149
150        let rows = instrument_named!(
151            sqlx::query(&query).fetch_all(&self.lake.db_pool),
152            "sql_select_list_partitions"
153        )
154        .await
155        .map_err(|e| DataFusionError::External(e.into()))?;
156        let rb = rows_to_record_batch(&rows).map_err(|e| DataFusionError::External(e.into()))?;
157
158        let source = MemorySourceConfig::try_new(
159            &[vec![rb]],
160            self.schema(),
161            projection.map(|v| v.to_owned()),
162        )?;
163        Ok(DataSourceExec::from_data_source(source))
164    }
165}
micromegas_analytics/lakehouse/list_partitions_table_function.rs

micromegas_analytics/lakehouse/
list_partitions_table_function.rs
