micromegas_analytics/lakehouse/
sql_batch_view.rs

1use super::{
2    batch_update::PartitionCreationStrategy,
3    dataframe_time_bounds::{DataFrameTimeBounds, NamedColumnsTimeBounds},
4    lakehouse_context::LakehouseContext,
5    materialized_view::MaterializedView,
6    merge::{MergeQueryResult, PartitionMerger, QueryMerger},
7    partition::Partition,
8    partition_cache::{NullPartitionProvider, PartitionCache},
9    partitioned_execution_plan::ScanOrdering,
10    query::make_session_context,
11    read_scope::CallerContext,
12    session_configurator::SessionConfigurator,
13    sql_partition_spec::fetch_sql_partition_spec,
14    view::{PartitionSpec, ScanSortColumn, View, ViewMetadata},
15    view_factory::ViewFactory,
16};
17use crate::{
18    record_batch_transformer::TrivialRecordBatchTransformer,
19    time::{TimeRange, datetime_to_scalar},
20};
21use anyhow::{Context, Result};
22use async_trait::async_trait;
23use chrono::{DateTime, TimeDelta, Utc};
24use datafusion::{
25    arrow::datatypes::Schema, common::TableReference, execution::runtime_env::RuntimeEnv,
26    prelude::*,
27};
28use micromegas_ingestion::data_lake_connection::DataLakeConnection;
29use micromegas_tracing::error;
30use std::hash::Hash;
31use std::hash::Hasher;
32use std::{hash::DefaultHasher, sync::Arc};
33
34/// The `SQLOptions` every plan-only build of a stored query text (extract/count/merge, at
35/// construction or validation time) must use instead of the unguarded `ctx.sql(...)`:
36/// `SQLOptions::new()`'s defaults are `allow_ddl`/`allow_dml`/`allow_statements: true`, so planning
37/// a stored query containing an embedded `CREATE EXTERNAL TABLE`/`DROP VIEW`/`COPY ... TO` would
38/// otherwise execute it once here and again on every daemon tick (`SqlPartitionSpec::write` plans
39/// `extract_query` the same way it was first planned). `pub(crate)` so `view_definition.rs`'s
40/// validation checks reuse the exact same guard.
41pub(crate) fn guarded_sql_options() -> SQLOptions {
42    SQLOptions::new()
43        .with_allow_ddl(false)
44        .with_allow_dml(false)
45        .with_allow_statements(false)
46}
47
48/// Builds the ascending `ScanSortColumn` list corresponding to a declared sort order.
49fn sort_order_as_scan_columns(columns: &[Arc<String>]) -> Vec<ScanSortColumn> {
50    columns
51        .iter()
52        .map(|column| ScanSortColumn {
53            column: column.clone(),
54            descending: false,
55        })
56        .collect()
57}
58
59/// SQL-defined view updated in batch
60#[derive(Debug)]
61pub struct SqlBatchView {
62    view_set_name: Arc<String>,
63    view_instance_id: Arc<String>,
64    min_event_time_column: Arc<String>,
65    max_event_time_column: Arc<String>,
66    count_src_query: Arc<String>,
67    extract_query: Arc<String>,
68    merge_partitions_query: Arc<String>,
69    schema: Arc<Schema>,
70    merger: Arc<dyn PartitionMerger>,
71    view_factory: Arc<ViewFactory>,
72    session_configurator: Arc<dyn SessionConfigurator>,
73    update_group: Option<i32>,
74    max_partition_delta_from_source: TimeDelta,
75    max_partition_delta_from_merge: TimeDelta,
76    /// The ascending sort order this view's partitions are declared to satisfy (see
77    /// `with_merge_sort_order`), if any.
78    sort_order: Option<Vec<Arc<String>>>,
79    /// A `QueryMerger` declaring `ScanOrdering::PerFile { columns: sort_order }`, used instead of
80    /// `merger` when every input to a merge certifies `sort_order`.
81    ordered_merger: Option<Arc<dyn PartitionMerger>>,
82}
83
84impl SqlBatchView {
85    #[expect(clippy::too_many_arguments)]
86    /// # Arguments
87    ///
88    /// * `runtime` - datafusion runtime
89    /// * `view_set_name` - name of the table
90    /// * `min_event_time_column` - min(column) should result in the first timestamp in a dataframe
91    /// * `max_event_time_column` - max(column) should result in the last timestamp in a dataframe
92    /// * `count_src_query` - used to count the rows of the underlying data to know if a cached partition is up to date
93    /// * `extract_query` - used to extract the source data into a cached partition
94    /// * `merge_partitions_query` - used to merge multiple partitions into a single one (and user queries which are one multiple partitions by default)
95    /// * `lake` - data lake
96    /// * `view_factory` - all views accessible to the `count_src_query`
97    /// * `session_configurator` - configurator for registering custom tables (e.g., JSON files)
98    /// * `update_group` - tells the daemon which view should be materialized and in what order
99    pub async fn new(
100        runtime: Arc<RuntimeEnv>,
101        view_set_name: Arc<String>,
102        min_event_time_column: Arc<String>,
103        max_event_time_column: Arc<String>,
104        count_src_query: Arc<String>,
105        extract_query: Arc<String>,
106        merge_partitions_query: Arc<String>,
107        lake: Arc<DataLakeConnection>,
108        view_factory: Arc<ViewFactory>,
109        session_configurator: Arc<dyn SessionConfigurator>,
110        update_group: Option<i32>,
111        max_partition_delta_from_source: TimeDelta,
112        max_partition_delta_from_merge: TimeDelta,
113    ) -> Result<Self> {
114        let null_part_provider = Arc::new(NullPartitionProvider {});
115        let lakehouse = Arc::new(LakehouseContext::new(lake.clone(), runtime.clone())?);
116        let ctx = make_session_context(
117            lakehouse,
118            null_part_provider,
119            None,
120            view_factory.clone(),
121            session_configurator.clone(),
122            CallerContext::maintenance(),
123        )
124        .await
125        .with_context(|| "make_session_context")?;
126        let now_str = Utc::now().to_rfc3339();
127        let sql = extract_query
128            .replace("{begin}", &now_str)
129            .replace("{end}", &now_str);
130        // Guarded: an embedded DDL/DML statement in extract_query must never execute here (see
131        // `guarded_sql_options`'s doc comment) -- this is what makes the DDL validator's own check
132        // 10 hold even though construction runs before validation.
133        let extracted_df = ctx.sql_with_options(&sql, guarded_sql_options()).await?;
134        let schema = extracted_df.schema().inner().clone();
135        let merge_query = Arc::new(merge_partitions_query.replace("{source}", "source"));
136        let merger: Arc<dyn PartitionMerger> = Arc::new(QueryMerger::new(
137            view_factory.clone(),
138            session_configurator.clone(),
139            schema.clone(),
140            merge_query,
141        ));
142
143        Ok(Self {
144            view_set_name,
145            view_instance_id: Arc::new(String::from("global")),
146            min_event_time_column,
147            max_event_time_column,
148            count_src_query,
149            extract_query,
150            merge_partitions_query,
151            schema,
152            merger,
153            view_factory,
154            session_configurator,
155            update_group,
156            max_partition_delta_from_source,
157            max_partition_delta_from_merge,
158            sort_order: None,
159            ordered_merger: None,
160        })
161    }
162
163    /// Declares that this view's partitions are internally sorted, ascending, by `columns` in
164    /// order. `columns` must be non-empty and every name must exist in this view's schema.
165    ///
166    /// This is a three-item view-author contract:
167    /// 1. Every declared sort column must appear among the merge query's `GROUP BY` keys (order
168    ///    within `GROUP BY` is irrelevant; extra keys degrade to `PartiallySorted`, not a blocking
169    ///    sort).
170    /// 2. Any enrichment join in the merge query must put the dimension table on the left and the
171    ///    ordered stream on the right -- `CollectLeft` buffers its left (build) input and inherits
172    ///    its output ordering from its right (probe) input, so the natural phrasing
173    ///    `(<ordered agg>) a LEFT JOIN dim d` reinstates a blocking sort.
174    /// 3. The merge query's aggregates must be composable over already-aggregated rows (e.g.
175    ///    `sum(count)`, not `count(*)`; no bare `avg` -- carry `sum` and `count` and divide at
176    ///    read time). `log_stats` is the in-repo model.
177    ///
178    /// Neither query needs an author-written `ORDER BY` any more: this builder forwards `columns`
179    /// to both sides, `QueryMerger` applies the sort to the merge query as a DataFusion
180    /// logical-plan node, and the extract path applies it to the extract query the same way (see
181    /// `sql_partition_spec::plan_sorted_extract`) -- no SQL text is derived or rewritten on either
182    /// side.
183    pub fn with_merge_sort_order(mut self, columns: Vec<Arc<String>>) -> Result<Self> {
184        if columns.is_empty() {
185            anyhow::bail!("with_merge_sort_order: columns must be non-empty");
186        }
187        for column in &columns {
188            self.schema.index_of(column.as_str()).with_context(|| {
189                format!(
190                    "with_merge_sort_order: column {column:?} not found in view {}'s schema",
191                    self.view_set_name
192                )
193            })?;
194        }
195        let scan_columns = sort_order_as_scan_columns(&columns);
196        let merge_query = Arc::new(self.merge_partitions_query.replace("{source}", "source"));
197        let ordered_merger: Arc<dyn PartitionMerger> = Arc::new(
198            QueryMerger::new(
199                self.view_factory.clone(),
200                self.session_configurator.clone(),
201                self.schema.clone(),
202                merge_query,
203            )
204            .with_merge_scan_ordering(ScanOrdering::PerFile {
205                columns: scan_columns,
206            }),
207        );
208        self.ordered_merger = Some(ordered_merger);
209        self.sort_order = Some(columns);
210        Ok(self)
211    }
212
213    /// The extract query as given to `SqlBatchView::new`, `{begin}`/`{end}` placeholders included.
214    pub fn get_extract_query(&self) -> Arc<String> {
215        self.extract_query.clone()
216    }
217
218    /// The merge query as given to `SqlBatchView::new`, `{source}` placeholder included.
219    pub fn get_merge_partitions_query(&self) -> Arc<String> {
220        self.merge_partitions_query.clone()
221    }
222
223    /// True when a merge sort order is declared and every partition in `partitions_to_merge`
224    /// certifies it (empty partitions certify vacuously -- see `Partition::certifies_sort_order`).
225    fn all_inputs_certify(&self, partitions_to_merge: &[Partition]) -> bool {
226        let Some(sort_order) = &self.sort_order else {
227            return false;
228        };
229        let columns = sort_order_as_scan_columns(sort_order);
230        partitions_to_merge
231            .iter()
232            .all(|p| p.certifies_sort_order(&columns))
233    }
234}
235
236#[async_trait]
237impl View for SqlBatchView {
238    fn get_view_set_name(&self) -> Arc<String> {
239        self.view_set_name.clone()
240    }
241
242    fn get_view_instance_id(&self) -> Arc<String> {
243        self.view_instance_id.clone()
244    }
245
246    async fn make_batch_partition_spec(
247        &self,
248        lakehouse: Arc<LakehouseContext>,
249        existing_partitions: Arc<PartitionCache>,
250        insert_range: TimeRange,
251    ) -> Result<Arc<dyn PartitionSpec>> {
252        let view_meta = ViewMetadata {
253            view_set_name: self.get_view_set_name(),
254            view_instance_id: self.get_view_instance_id(),
255            file_schema_hash: self.get_file_schema_hash(),
256        };
257        let partitions_in_range = Arc::new(existing_partitions.filter_insert_range(insert_range));
258        let ctx = make_session_context(
259            lakehouse,
260            partitions_in_range.clone(),
261            None,
262            self.view_factory.clone(),
263            self.session_configurator.clone(),
264            CallerContext::maintenance(),
265        )
266        .await
267        .with_context(|| "make_session_context")?;
268
269        let count_src_sql = self
270            .count_src_query
271            .replace("{begin}", &insert_range.begin.to_rfc3339())
272            .replace("{end}", &insert_range.end.to_rfc3339());
273
274        let extract_sql = self
275            .extract_query
276            .replace("{begin}", &insert_range.begin.to_rfc3339())
277            .replace("{end}", &insert_range.end.to_rfc3339());
278
279        Ok(Arc::new(
280            fetch_sql_partition_spec(
281                ctx,
282                Arc::new(TrivialRecordBatchTransformer {}),
283                self.get_time_bounds(),
284                self.schema.clone(),
285                count_src_sql,
286                extract_sql,
287                view_meta,
288                insert_range,
289                self.sort_order
290                    .as_ref()
291                    .map(|columns| columns.iter().map(|c| c.to_string()).collect()),
292            )
293            .await
294            .with_context(|| "fetch_sql_partition_spec")?,
295        ))
296    }
297
298    fn get_file_schema_hash(&self) -> Vec<u8> {
299        let mut hasher = DefaultHasher::new();
300        self.schema.hash(&mut hasher);
301        hasher.finish().to_le_bytes().to_vec()
302    }
303
304    fn get_file_schema(&self) -> Arc<Schema> {
305        self.schema.clone()
306    }
307
308    async fn jit_update(
309        &self,
310        _lakehouse: Arc<LakehouseContext>,
311        _query_range: Option<TimeRange>,
312    ) -> Result<()> {
313        Ok(())
314    }
315
316    fn make_time_filter(&self, begin: DateTime<Utc>, end: DateTime<Utc>) -> Result<Vec<Expr>> {
317        Ok(vec![
318            col(&*self.min_event_time_column).lt_eq(lit(datetime_to_scalar(end))),
319            col(&*self.max_event_time_column).gt_eq(lit(datetime_to_scalar(begin))),
320        ])
321    }
322
323    fn get_time_bounds(&self) -> Arc<dyn DataFrameTimeBounds> {
324        Arc::new(NamedColumnsTimeBounds::new(
325            self.min_event_time_column.clone(),
326            self.max_event_time_column.clone(),
327        ))
328    }
329
330    async fn register_table(&self, ctx: &SessionContext, table: MaterializedView) -> Result<()> {
331        let view_name = self.get_view_set_name().to_string();
332        let partitions_table_name = format!("__{view_name}__partitions");
333        ctx.register_table(
334            TableReference::Bare {
335                table: partitions_table_name.clone().into(),
336            },
337            Arc::new(table),
338        )?;
339        let df = ctx
340            .sql(
341                &self
342                    .merge_partitions_query
343                    .replace("{source}", &partitions_table_name),
344            )
345            .await?;
346        ctx.register_table(
347            TableReference::Bare {
348                table: view_name.into(),
349            },
350            df.into_view(),
351        )?;
352        Ok(())
353    }
354
355    async fn merge_partitions(
356        &self,
357        lakehouse: Arc<LakehouseContext>,
358        partitions_to_merge: Arc<Vec<Partition>>,
359        partitions_all_views: Arc<PartitionCache>,
360        insert_range: TimeRange,
361    ) -> Result<MergeQueryResult> {
362        // An all-empty source scans as an EmptyExec, whose SortExec is never elided -- taking the
363        // ordered path there would trip the plan-shape check's memory-regression warning on every
364        // quiet-day retry. So the ordered path additionally requires at least one non-empty input.
365        let any_non_empty = partitions_to_merge.iter().any(|p| !p.is_empty());
366        let merger = match &self.ordered_merger {
367            Some(ordered_merger)
368                if any_non_empty && self.all_inputs_certify(&partitions_to_merge) =>
369            {
370                ordered_merger
371            }
372            _ => &self.merger,
373        };
374        let res = merger
375            .execute_merge_query(
376                lakehouse,
377                partitions_to_merge,
378                partitions_all_views,
379                insert_range,
380            )
381            .await;
382        if let Err(e) = &res {
383            error!("{e:?}");
384        }
385        res
386    }
387
388    fn get_merged_partition_sort_order(
389        &self,
390        partitions_to_merge: &[Partition],
391    ) -> Option<Vec<String>> {
392        if self.ordered_merger.is_some() && self.all_inputs_certify(partitions_to_merge) {
393            self.sort_order
394                .as_ref()
395                .map(|columns| columns.iter().map(|c| c.to_string()).collect())
396        } else {
397            None
398        }
399    }
400
401    fn get_scan_output_ordering(&self) -> ScanOrdering {
402        match &self.sort_order {
403            Some(columns) => ScanOrdering::PerFile {
404                columns: sort_order_as_scan_columns(columns),
405            },
406            None => ScanOrdering::Unordered,
407        }
408    }
409
410    fn get_update_group(&self) -> Option<i32> {
411        self.update_group
412    }
413
414    fn get_max_partition_time_delta(&self, strategy: &PartitionCreationStrategy) -> TimeDelta {
415        match strategy {
416            PartitionCreationStrategy::Abort | PartitionCreationStrategy::CreateFromSource => {
417                self.max_partition_delta_from_source
418            }
419            PartitionCreationStrategy::MergeExisting(_partitions) => {
420                self.max_partition_delta_from_merge
421            }
422        }
423    }
424}
micromegas_analytics/lakehouse/sql_batch_view.rs

micromegas_analytics/lakehouse/
sql_batch_view.rs
