micromegas_analytics/lakehouse/
sql_partition_spec.rs

1use super::{
2    dataframe_time_bounds::DataFrameTimeBounds,
3    partitioned_execution_plan::{assert_ordering_satisfied, assert_single_partition, sort_exprs},
4    view::{PartitionSpec, ScanSortColumn, ViewMetadata},
5    write_partition::write_partition_from_rows,
6};
7use crate::{
8    dfext::typed_column::typed_column_by_name,
9    lakehouse::write_partition::{PartitionRowSet, RetireMatch},
10    record_batch_transformer::RecordBatchTransformer,
11    response_writer::Logger,
12    time::TimeRange,
13};
14use anyhow::{Context, Result};
15use async_trait::async_trait;
16use datafusion::{
17    arrow::{
18        array::{Int64Array, RecordBatch},
19        datatypes::Schema,
20    },
21    execution::{SendableRecordBatchStream, TaskContext},
22    physical_plan::{ExecutionPlan, execute_stream},
23    prelude::*,
24};
25use futures::StreamExt;
26use micromegas_ingestion::data_lake_connection::DataLakeConnection;
27use micromegas_tracing::prelude::*;
28use std::sync::Arc;
29
30/// A `PartitionSpec` implementation for SQL-defined partitions.
31pub struct SqlPartitionSpec {
32    ctx: SessionContext,
33    transformer: Arc<dyn RecordBatchTransformer>,
34    compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
35    schema: Arc<Schema>,
36    extract_query: String,
37    view_metadata: ViewMetadata,
38    insert_range: TimeRange,
39    record_count: i64,
40    /// The sort guarantee to record on the fresh partition this extract query writes, if any (see
41    /// `SqlBatchView::with_merge_sort_order`). When set, `write` applies it to the extract
42    /// query's `DataFrame` as a `DataFrame::sort` before executing it (see
43    /// `plan_sorted_extract`), so the guarantee it records always matches what was actually
44    /// written.
45    sort_order: Option<Vec<String>>,
46}
47
48impl SqlPartitionSpec {
49    #[expect(clippy::too_many_arguments)]
50    pub fn new(
51        ctx: SessionContext,
52        transformer: Arc<dyn RecordBatchTransformer>,
53        compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
54        schema: Arc<Schema>,
55        extract_query: String,
56        view_metadata: ViewMetadata,
57        insert_range: TimeRange,
58        record_count: i64,
59        sort_order: Option<Vec<String>>,
60    ) -> Self {
61        Self {
62            ctx,
63            transformer,
64            compute_time_bounds,
65            schema,
66            extract_query,
67            view_metadata,
68            insert_range,
69            record_count,
70            sort_order,
71        }
72    }
73
74    /// Converts the declared `sort_order` (if any) to `ScanSortColumn`s and delegates to
75    /// `plan_sorted_extract`, which applies it to the extract query's `DataFrame` before building
76    /// the physical plan -- the same apply-then-plan discipline as
77    /// `QueryMerger::execute_sorted_merge` -- then executes the plan it returns. The undeclared
78    /// path (`sort_order: None`) plans and executes the query unsorted.
79    async fn execute_extract_query(&self, df: DataFrame) -> Result<SendableRecordBatchStream> {
80        let columns: Option<Vec<ScanSortColumn>> = self.sort_order.as_ref().map(|sort_order| {
81            sort_order
82                .iter()
83                .map(|c| ScanSortColumn {
84                    column: Arc::new(c.clone()),
85                    descending: false,
86                })
87                .collect()
88        });
89        let subject = format!("extract query for {}", self.view_metadata.view_set_name);
90        let (plan, task_ctx) =
91            plan_sorted_extract(df, columns.as_deref(), &subject, self.insert_range).await?;
92        execute_stream(plan, task_ctx).with_context(|| "executing extract query plan")
93    }
94}
95
96/// Applies `sort_order` (if declared) to `df` as a `DataFrame::sort` over its `ScanSortColumn`s,
97/// builds the physical plan, and -- only when a sort order was declared -- runs
98/// `assert_single_partition` and `assert_ordering_satisfied` against it, mirroring
99/// `QueryMerger::execute_sorted_merge`'s apply-then-plan sequence (`merge.rs:227-270`; its four
100/// merge-specific optimizer settings have no counterpart here). With `sort_order: None` the plan
101/// is built unsorted and neither assertion runs. `subject` names the query for the assertions'
102/// error messages (e.g. `format!("extract query for {view}")`); `insert_range` is the range being
103/// written. Returns the physical plan and the `TaskContext` the caller needs to execute it --
104/// callers are left to `execute_stream` it themselves, since only they know what to do with the
105/// resulting stream.
106pub async fn plan_sorted_extract(
107    df: DataFrame,
108    sort_order: Option<&[ScanSortColumn]>,
109    subject: &str,
110    insert_range: TimeRange,
111) -> Result<(Arc<dyn ExecutionPlan>, Arc<TaskContext>)> {
112    let Some(columns) = sort_order else {
113        let task_ctx = Arc::new(df.task_ctx());
114        let plan = df
115            .create_physical_plan()
116            .await
117            .with_context(|| "creating physical plan for extract query")?;
118        return Ok((plan, task_ctx));
119    };
120
121    let df = df.sort(sort_exprs(columns))?;
122    let task_ctx = Arc::new(df.task_ctx());
123    let plan = df
124        .create_physical_plan()
125        .await
126        .with_context(|| "creating physical plan for extract query")?;
127
128    assert_single_partition(
129        &plan,
130        subject,
131        insert_range,
132        "a declared sort_order requires a single-partition, globally-ordered output.",
133    )?;
134
135    assert_ordering_satisfied(
136        &plan,
137        columns,
138        "extract-query",
139        subject,
140        insert_range,
141        &format!(
142            "the declared sort_order {columns:?} even after applying it as a DataFrame::sort; \
143             this indicates the physical plan silently dropped or reordered the applied sort, \
144             not an author mistake."
145        ),
146    )?;
147
148    Ok((plan, task_ctx))
149}
150
151impl std::fmt::Debug for SqlPartitionSpec {
152    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153        write!(f, "SqlPartitionSpec")
154    }
155}
156
157#[async_trait]
158impl PartitionSpec for SqlPartitionSpec {
159    fn is_empty(&self) -> bool {
160        self.record_count < 1
161    }
162
163    fn get_source_data_hash(&self) -> Vec<u8> {
164        self.record_count.to_le_bytes().to_vec()
165    }
166
167    async fn write(&self, lake: Arc<DataLakeConnection>, logger: Arc<dyn Logger>) -> Result<()> {
168        // Allow empty record_count - write_partition_from_rows will create
169        // an empty partition record if no data is sent through the channel
170        let desc = format!(
171            "[{}, {}] {} {}",
172            self.view_metadata.view_set_name,
173            self.view_metadata.view_instance_id,
174            self.insert_range.begin.to_rfc3339(),
175            self.insert_range.end.to_rfc3339()
176        );
177        logger.write_log_entry(format!("writing {desc}")).await?;
178        let df = self.ctx.sql(&self.extract_query).await?;
179        let mut stream = self.execute_extract_query(df).await?;
180
181        let (tx, rx) = tokio::sync::mpsc::channel(1);
182        let join_handle = spawn_with_context(write_partition_from_rows(
183            lake.clone(),
184            self.view_metadata.clone(),
185            self.schema.clone(),
186            self.insert_range,
187            self.get_source_data_hash(),
188            self.sort_order.clone(),
189            RetireMatch::Containment,
190            Vec::new(),
191            rx,
192            logger.clone(),
193        ));
194
195        while let Some(rb_res) = stream.next().await {
196            let rb = self.transformer.transform(rb_res?).await?;
197            let event_time_range = self
198                .compute_time_bounds
199                .get_time_bounds(self.ctx.read_batch(rb.clone())?)
200                .await?;
201            tx.send(Ok(PartitionRowSet::new(event_time_range, rb, None)))
202                .await?;
203        }
204        drop(tx);
205        join_handle.await??;
206        Ok(())
207    }
208}
209
210/// Fetches a `SqlPartitionSpec` by executing a count query and an extract query. `sort_order`
211/// declares the ordering the resulting partition's rows are guaranteed to satisfy (see
212/// `SqlBatchView::with_merge_sort_order`); pass `None` for views that make no such guarantee.
213#[expect(clippy::too_many_arguments)]
214pub async fn fetch_sql_partition_spec(
215    ctx: SessionContext,
216    transformer: Arc<dyn RecordBatchTransformer>,
217    compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
218    schema: Arc<Schema>,
219    count_src_sql: String,
220    extract_query: String,
221    view_metadata: ViewMetadata,
222    insert_range: TimeRange,
223    sort_order: Option<Vec<String>>,
224) -> Result<SqlPartitionSpec> {
225    let df = ctx.sql(&count_src_sql).await?;
226    let batches: Vec<RecordBatch> = df.collect().await?;
227    if batches.len() != 1 {
228        anyhow::bail!("fetch_sql_partition_spec: query should return a single batch");
229    }
230    let rb = &batches[0];
231    let count_column: &Int64Array = typed_column_by_name(rb, "count")?;
232    if count_column.len() != 1 {
233        anyhow::bail!("fetch_sql_partition_spec: query should return a single row");
234    }
235    let count = count_column.value(0);
236    if count > 0 {
237        trace!(
238            "fetch_sql_partition_spec for view {}, count={count}",
239            &*view_metadata.view_set_name
240        );
241    }
242    Ok(SqlPartitionSpec::new(
243        ctx,
244        transformer,
245        compute_time_bounds,
246        schema,
247        extract_query,
248        view_metadata,
249        insert_range,
250        count,
251        sort_order,
252    ))
253}
micromegas_analytics/lakehouse/sql_partition_spec.rs

micromegas_analytics/lakehouse/
sql_partition_spec.rs
