micromegas_analytics/lakehouse/
list_partitions_table_function.rs

1use super::audience_guard::{AudienceGuard, IdKind};
2use super::read_scope::ReadScope;
3use crate::sql_arrow_bridge::rows_to_record_batch;
4use async_trait::async_trait;
5use datafusion::arrow::array::RecordBatch;
6use datafusion::arrow::datatypes::DataType;
7use datafusion::arrow::datatypes::Field;
8use datafusion::arrow::datatypes::Schema;
9use datafusion::arrow::datatypes::SchemaRef;
10use datafusion::arrow::datatypes::TimeUnit;
11use datafusion::catalog::Session;
12use datafusion::catalog::TableFunctionArgs;
13use datafusion::catalog::TableFunctionImpl;
14use datafusion::catalog::TableProvider;
15use datafusion::datasource::TableType;
16use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
17use datafusion::error::DataFusionError;
18use datafusion::physical_plan::ExecutionPlan;
19use datafusion::prelude::Expr;
20use micromegas_ingestion::data_lake_connection::DataLakeConnection;
21use micromegas_tracing::prelude::*;
22use sqlx::Row;
23use std::collections::HashSet;
24use std::sync::Arc;
25use uuid::Uuid;
26
27/// The 15-column projection shared by every `lakehouse_partitions` query in [`ListPartitionsTableProvider::scan`],
28/// kept in sync with [`ListPartitionsTableProvider::schema`]'s field order.
29const SELECT_PARTITION_COLUMNS: &str = "SELECT view_set_name,
30             view_instance_id,
31             begin_insert_time,
32             end_insert_time,
33             min_event_time,
34             max_event_time,
35             updated,
36             file_path,
37             file_size,
38             file_schema_hash,
39             source_data_hash,
40             num_rows,
41             partition_format_version,
42             sort_order,
43             max_sort_key_time";
44
45/// A DataFusion `TableFunctionImpl` for listing lakehouse partitions.
46#[derive(Debug)]
47pub struct ListPartitionsTableFunction {
48    lake: Arc<DataLakeConnection>,
49    guard: Arc<AudienceGuard>,
50}
51
52impl ListPartitionsTableFunction {
53    pub fn new(lake: Arc<DataLakeConnection>, guard: Arc<AudienceGuard>) -> Self {
54        Self { lake, guard }
55    }
56}
57
58impl TableFunctionImpl for ListPartitionsTableFunction {
59    fn call_with_args(
60        &self,
61        _args: TableFunctionArgs,
62    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
63        Ok(Arc::new(ListPartitionsTableProvider {
64            lake: self.lake.clone(),
65            guard: self.guard.clone(),
66        }))
67    }
68}
69
70/// A DataFusion `TableProvider` for listing lakehouse partitions.
71#[derive(Debug)]
72pub struct ListPartitionsTableProvider {
73    pub lake: Arc<DataLakeConnection>,
74    guard: Arc<AudienceGuard>,
75}
76
77#[async_trait]
78impl TableProvider for ListPartitionsTableProvider {
79    fn schema(&self) -> SchemaRef {
80        Arc::new(Schema::new(vec![
81            Field::new("view_set_name", DataType::Utf8, false),
82            Field::new("view_instance_id", DataType::Utf8, false),
83            Field::new(
84                "begin_insert_time",
85                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
86                false,
87            ),
88            Field::new(
89                "end_insert_time",
90                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
91                false,
92            ),
93            Field::new(
94                "min_event_time",
95                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
96                true,
97            ),
98            Field::new(
99                "max_event_time",
100                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
101                true,
102            ),
103            Field::new(
104                "updated",
105                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
106                false,
107            ),
108            Field::new("file_path", DataType::Utf8, true),
109            Field::new("file_size", DataType::Int64, false),
110            Field::new("file_schema_hash", DataType::Binary, false),
111            Field::new("source_data_hash", DataType::Binary, false),
112            Field::new("num_rows", DataType::Int64, false),
113            Field::new("partition_format_version", DataType::Int32, false),
114            Field::new(
115                "sort_order",
116                DataType::List(Arc::new(Field::new("tag", DataType::Utf8, false))),
117                true,
118            ),
119            Field::new(
120                "max_sort_key_time",
121                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
122                true,
123            ),
124        ]))
125    }
126
127    fn table_type(&self) -> TableType {
128        TableType::Temporary
129    }
130
131    async fn scan(
132        &self,
133        _state: &dyn Session,
134        projection: Option<&Vec<usize>>,
135        _filters: &[Expr],
136        limit: Option<usize>,
137    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
138        // The call-level guard: `ReadScope::All` keeps the unrestricted path unchanged,
139        // including the `LIMIT` pushdown. A restricted caller instead fetches every row unlimited
140        // and row-filters, with `filter_rows` itself enforcing `limit` by stopping as soon as
141        // enough matching rows are kept -- filtering after a pushed-down limit would return fewer
142        // rows than asked for while more matching rows exist, silently wrong.
143        let restricted = *self.guard.read_scope() != ReadScope::All;
144
145        // Build query with optional LIMIT clause pushed down to PostgreSQL, but only when
146        // unrestricted. DataFusion only pushes the limit when it's safe to do so (i.e., when
147        // there are no WHERE clauses that could filter rows). When filters are present,
148        // DataFusion passes limit=None and applies the limit after filtering.
149        // Important: DataFusion trusts us to apply the limit - if we ignore it, too many rows
150        // will be returned to the client. A restricted caller must never push the limit down --
151        // see the comment on `filter_rows` for why (deliberately unbounded fetch below).
152        let limit_clause = match (restricted, limit) {
153            (false, Some(n)) => format!("\n             LIMIT {n}"),
154            _ => String::new(),
155        };
156        // Deliberately unbounded when restricted (no `LIMIT`/`OFFSET` paging here): paginating
157        // this fetch would need either (a) holding a connection/transaction across the paging
158        // loop while `filter_rows` re-enters this same `db_pool` for every audience-resolution
159        // chunk -- self-deadlocking once concurrent restricted scans exceed the pool's connection
160        // count -- or (b) a stable, indexable ordering key to page on, which
161        // `lakehouse_partitions` doesn't have (no primary key, `file_path` is nullable, and the
162        // only exclusion constraint covers just overlap within one `(view_set_name,
163        // view_instance_id, file_schema_hash)` group). A single unbounded fetch avoids both.
164        let query = format!("{SELECT_PARTITION_COLUMNS} FROM lakehouse_partitions{limit_clause};");
165
166        let rows = instrument_named!(
167            sqlx::query(&query).fetch_all(&self.lake.db_pool),
168            "sql_select_list_partitions"
169        )
170        .await
171        .map_err(|e| DataFusionError::External(e.into()))?;
172
173        let rows = if restricted {
174            self.filter_rows(rows, limit).await?
175        } else {
176            rows
177        };
178        // `rows_to_record_batch` maps an empty slice to a **zero-field** empty batch
179        // (`make_empty_record_batch`), which doesn't match this provider's 15-column schema and
180        // fails once `MemoryStream` projects it. Build the empty batch directly from this
181        // provider's own schema instead -- this is the normal state of both a fresh deployment
182        // with no partitions materialized yet, and a `ReadScope::Audiences` caller with no
183        // readable partitions. Note `filter_rows` already stops as soon as `limit` matching rows
184        // have been kept (see its doc comment), so `rows.len() <= limit.unwrap_or(usize::MAX)`
185        // always holds here -- no further truncation needed.
186        let rb = if rows.is_empty() {
187            RecordBatch::new_empty(self.schema())
188        } else {
189            rows_to_record_batch(&rows).map_err(|e| DataFusionError::External(e.into()))?
190        };
191
192        let source = MemorySourceConfig::try_new(
193            &[vec![rb]],
194            self.schema(),
195            projection.map(|v| v.to_owned()),
196        )?;
197        Ok(DataSourceExec::from_data_source(source))
198    }
199}
200
201impl ListPartitionsTableProvider {
202    /// Bounds only how many rows' worth of ids a single `readable_ids` round trip resolves --
203    /// it does NOT bound the row fetch from Postgres, which is a single unbounded `fetch_all`
204    /// (see the comment on that query in [`ListPartitionsTableProvider::scan`]). Keeps a
205    /// restricted caller's `LIMIT n` from resolving every distinct `view_instance_id` in
206    /// `lakehouse_partitions` up front: already-fetched rows are processed in chunks of this
207    /// size, stopping as soon as `limit` matching rows have been kept, so a small `LIMIT` costs
208    /// at most a handful of audience-resolution round trips instead of one unbounded batch.
209    const RESOLVE_CHUNK_ROWS: usize = 1_000;
210
211    /// Row filtering for a `ReadScope::Audiences` caller: a `view_instance_id` that
212    /// parses as a `Uuid` is kept iff it resolves (as `IdKind::ProcessOrStream`) to a readable
213    /// audience; the literal `'global'` is kept iff `AudienceGuard::global_rows_visible` says so
214    /// for that row's `view_set_name`; anything else is dropped (fail-closed -- nothing produces
215    /// such a value today). Filters the `sqlx` row vector *before* `rows_to_record_batch` rather
216    /// than the built `RecordBatch` after: simpler than a `take` kernel over 15 columns, and
217    /// leaves the schema construction untouched. Processes `rows` in bounded chunks
218    /// ([`Self::RESOLVE_CHUNK_ROWS`]) and returns as soon as `limit` matching rows have been kept.
219    async fn filter_rows(
220        &self,
221        rows: Vec<sqlx::postgres::PgRow>,
222        limit: Option<usize>,
223    ) -> datafusion::error::Result<Vec<sqlx::postgres::PgRow>> {
224        let mut kept = Vec::new();
225        let mut rows_iter = rows.into_iter();
226        loop {
227            let chunk: Vec<_> = (&mut rows_iter).take(Self::RESOLVE_CHUNK_ROWS).collect();
228            if chunk.is_empty() {
229                break;
230            }
231            let mut candidate_ids: Vec<Uuid> = Vec::with_capacity(chunk.len());
232            let mut seen_ids: HashSet<Uuid> = HashSet::with_capacity(chunk.len());
233            let mut row_ids: Vec<Option<Uuid>> = Vec::with_capacity(chunk.len());
234            for row in &chunk {
235                let view_instance_id: &str = row
236                    .try_get("view_instance_id")
237                    .map_err(|e| DataFusionError::External(e.into()))?;
238                let id = Uuid::parse_str(view_instance_id).ok();
239                if let Some(id) = id
240                    && seen_ids.insert(id)
241                {
242                    candidate_ids.push(id);
243                }
244                row_ids.push(id);
245            }
246            let readable = self
247                .guard
248                .readable_ids(&candidate_ids, IdKind::ProcessOrStream)
249                .await?;
250            for (row, id) in chunk.into_iter().zip(row_ids) {
251                let keep = match id {
252                    Some(uuid) => readable.contains(&uuid),
253                    None => {
254                        let view_instance_id: &str = row
255                            .try_get("view_instance_id")
256                            .map_err(|e| DataFusionError::External(e.into()))?;
257                        let view_set_name: &str = row
258                            .try_get("view_set_name")
259                            .map_err(|e| DataFusionError::External(e.into()))?;
260                        view_instance_id == "global"
261                            && self.guard.global_rows_visible(view_set_name)
262                    }
263                };
264                if keep {
265                    kept.push(row);
266                    if let Some(n) = limit
267                        && kept.len() >= n
268                    {
269                        return Ok(kept);
270                    }
271                }
272            }
273        }
274        Ok(kept)
275    }
276}
micromegas_analytics/lakehouse/list_partitions_table_function.rs

micromegas_analytics/lakehouse/
list_partitions_table_function.rs
