micromegas_analytics/lakehouse/
metadata_partition_spec.rs

1use super::{
2    dataframe_time_bounds::DataFrameTimeBounds,
3    view::{PartitionSpec, ViewMetadata},
4};
5use crate::{
6    lakehouse::write_partition::{PartitionRowSet, RetireMatch, write_partition_from_rows},
7    response_writer::Logger,
8    sql_arrow_bridge::rows_to_record_batch,
9    time::TimeRange,
10};
11use anyhow::{Context, Result};
12use async_trait::async_trait;
13use datafusion::{
14    arrow::{compute::cast, datatypes::Schema, record_batch::RecordBatch},
15    prelude::*,
16};
17use futures::TryStreamExt;
18use micromegas_ingestion::data_lake_connection::DataLakeConnection;
19use micromegas_tracing::prelude::*;
20use sqlx::{Row, postgres::PgRow};
21use std::sync::Arc;
22use tokio::sync::mpsc::Sender;
23
24/// Flush threshold on the estimated byte size of the pending chunk -- bounds peak memory to one
25/// ~8 MB chunk, not one day's worth of Postgres rows. Byte-based like the Parquet writer's own
26/// 100 MB flush (`write_partition.rs`), because a row-count threshold bounds nothing when a few
27/// rows carry MB-sized properties/objects_metadata payloads. Deliberately the only flush metric.
28const SOURCE_BYTES_PER_BATCH: usize = 8 * 1024 * 1024;
29
30#[derive(Debug)]
31pub struct MetadataPartitionSpec {
32    pub view_metadata: ViewMetadata,
33    pub schema: Arc<Schema>,
34    pub insert_range: TimeRange,
35    pub record_count: i64,
36    pub data_sql: Arc<String>,
37    /// Bound as `$3` when `data_sql` runs -- the deployment's `MICROMEGAS_DEFAULT_AUDIENCE`,
38    /// which `BlocksView`'s `COALESCE` resolves a never-stamped process's missing
39    /// audience to. `None` for a `data_sql` that references only `$1`/`$2`; `BlocksView` is the
40    /// only view that uses this module, so in practice it is always `Some`. Note the separate
41    /// `source_count_query` deliberately does **not** get this bind -- it has no `$3`.
42    pub default_audience: Option<Arc<str>>,
43    pub compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
44    /// The sort guarantee this partition's rows will carry, per the caller's `data_sql`'s
45    /// `ORDER BY` (e.g. `Some(["insert_time"])` for `BlocksView`). Recorded on `Partition` as-is.
46    pub sort_order: Option<Vec<String>>,
47    /// The unfiltered row count over the same join and range as `source_count_query`, read from
48    /// that query's `unfiltered` column when present -- `Some` only for `BlocksView`, `None` for
49    /// every other caller of this module. [`Self::write`] combines it with `record_count` to log
50    /// how many rows the audience-mismatch predicate excluded; the check lives there rather than
51    /// in `make_batch_partition_spec` so it only fires for partitions actually written, not every
52    /// scheduled maintenance pass.
53    pub unfiltered_count: Option<i64>,
54}
55
56/// `unfiltered - kept`, clamped at zero -- the count of rows the audience-mismatch predicate
57/// excluded from this partition's materialization. Pulled out as a pure function so the
58/// arithmetic behind [`MetadataPartitionSpec::write`]'s `warn!`/`imetric!` pair is unit-testable
59/// directly. The clamp guards `kept > unfiltered`, which the single atomic count query behind
60/// both values should never produce.
61pub fn mismatch_excluded_count(unfiltered: i64, kept: i64) -> i64 {
62    (unfiltered - kept).max(0)
63}
64
65#[expect(clippy::too_many_arguments)]
66pub async fn fetch_metadata_partition_spec(
67    pool: &sqlx::PgPool,
68    source_count_query: &str,
69    data_sql: Arc<String>,
70    view_metadata: ViewMetadata,
71    schema: Arc<Schema>,
72    insert_range: TimeRange,
73    compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
74    sort_order: Option<Vec<String>>,
75    default_audience: Option<Arc<str>>,
76) -> Result<MetadataPartitionSpec> {
77    //todo: extract this query to allow join (instead of source_table)
78    let row = instrument_named!(
79        sqlx::query(source_count_query)
80            .bind(insert_range.begin)
81            .bind(insert_range.end)
82            .fetch_one(pool),
83        "sql_select_source_count"
84    )
85    .await
86    .with_context(|| "select count source metadata")?;
87    // `unfiltered` is only present when `source_count_query` folds it in (`BlocksView`); every
88    // other caller's query has no such column, so this is `None` for them rather than an error.
89    let unfiltered_count: Option<i64> = row.try_get("unfiltered").ok();
90    Ok(MetadataPartitionSpec {
91        view_metadata,
92        schema,
93        insert_range,
94        record_count: row.try_get("count").with_context(|| "reading count")?,
95        data_sql,
96        default_audience,
97        compute_time_bounds,
98        sort_order,
99        unfiltered_count,
100    })
101}
102
103/// Estimates a row's payload size by summing its raw column value byte lengths, counting `NULL`
104/// and any non-byte-backed value as 0. This deliberately tracks the JSONB/binary columns
105/// (`properties`, `objects_metadata`, `dependencies_metadata`) that dominate blocks-view row
106/// width -- an allocator-exact footprint is not needed, only a flush-decision estimate.
107fn estimate_row_bytes(row: &PgRow) -> usize {
108    let mut total = 0usize;
109    for i in 0..row.len() {
110        if let Ok(raw) = row.try_get_raw(i)
111            && let Ok(bytes) = raw.as_bytes()
112        {
113            total += bytes.len();
114        }
115    }
116    total
117}
118
119/// Aligns a Postgres-inferred batch to the declared file schema's column *types*, casting column
120/// by column where they differ. `sql_arrow_bridge`'s mapping is keyed on the Postgres type name
121/// alone, so it cannot know that a `TEXT` column is declared as something narrower downstream:
122/// `blocks.audience` is `Dictionary(Int32, Utf8)` in the file schema (matching every other view's
123/// audience column) but arrives as plain `Utf8`.
124///
125/// Positional, matching how the parquet writer zips declared fields against batch columns with no
126/// name check (see `write_partition::check_non_nullable_columns`). Nullability and field metadata
127/// come from the *batch*, never from the declared schema, so this stays a pure type alignment and
128/// leaves the `NOT NULL` verdict to the write path's own guard.
129pub fn cast_to_file_schema(batch: RecordBatch, file_schema: &Schema) -> Result<RecordBatch> {
130    let batch_schema = batch.schema();
131    let mut fields = Vec::with_capacity(batch_schema.fields().len());
132    let mut columns = Vec::with_capacity(batch.num_columns());
133    let mut any_cast = false;
134    for (i, field) in batch_schema.fields().iter().enumerate() {
135        let column = batch.column(i);
136        match file_schema.fields().get(i) {
137            Some(declared) if declared.data_type() != field.data_type() => {
138                columns.push(cast(column, declared.data_type()).with_context(|| {
139                    format!(
140                        "casting column {i} ({}) from {:?} to the declared {:?}",
141                        field.name(),
142                        field.data_type(),
143                        declared.data_type()
144                    )
145                })?);
146                fields.push(Arc::new(
147                    field
148                        .as_ref()
149                        .clone()
150                        .with_data_type(declared.data_type().clone()),
151                ));
152                any_cast = true;
153            }
154            _ => {
155                columns.push(column.clone());
156                fields.push(field.clone());
157            }
158        }
159    }
160    if !any_cast {
161        return Ok(batch);
162    }
163    RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
164        .with_context(|| "rebuilding the batch after casting to the declared file schema")
165}
166
167/// Converts the accumulated chunk to a `RecordBatch`, computes its event-time bounds, and sends
168/// it as a `PartitionRowSet`. Clears `chunk` in place for reuse by the next flush.
169async fn flush_chunk(
170    chunk: &mut Vec<PgRow>,
171    ctx: &SessionContext,
172    file_schema: &Schema,
173    compute_time_bounds: &Arc<dyn DataFrameTimeBounds>,
174    tx: &Sender<Result<PartitionRowSet, anyhow::Error>>,
175) -> Result<()> {
176    let record_batch = cast_to_file_schema(
177        rows_to_record_batch(chunk).with_context(|| "converting rows to record batch")?,
178        file_schema,
179    )?;
180    chunk.clear();
181    let event_time_range = compute_time_bounds
182        .get_time_bounds(
183            ctx.read_batch(record_batch.clone())
184                .with_context(|| "read_batch")?,
185        )
186        .await?;
187    tx.send(Ok(PartitionRowSet::new(
188        event_time_range,
189        record_batch,
190        None,
191    )))
192    .await
193    .with_context(|| "sending partition row set")?;
194    Ok(())
195}
196
197#[async_trait]
198impl PartitionSpec for MetadataPartitionSpec {
199    fn is_empty(&self) -> bool {
200        self.record_count < 1
201    }
202
203    fn get_source_data_hash(&self) -> Vec<u8> {
204        self.record_count.to_le_bytes().to_vec()
205    }
206
207    async fn write(&self, lake: Arc<DataLakeConnection>, logger: Arc<dyn Logger>) -> Result<()> {
208        // Allow empty record_count - write_partition_from_rows will create
209        // an empty partition record if no data is sent through the channel
210        let desc = format!(
211            "[{}, {}] {} {}",
212            self.view_metadata.view_set_name,
213            self.view_metadata.view_instance_id,
214            self.insert_range.begin.to_rfc3339(),
215            self.insert_range.end.to_rfc3339()
216        );
217        logger.write_log_entry(format!("writing {desc}")).await?;
218
219        // Logged/metered only here -- once a partition is actually about to be written, never on
220        // a scheduled pass that decides nothing needs writing (see `unfiltered_count`'s doc
221        // comment). The count itself was already fetched alongside `record_count` back in
222        // `fetch_metadata_partition_spec`, so this is pure arithmetic, not a second query.
223        if let Some(unfiltered) = self.unfiltered_count {
224            let excluded = mismatch_excluded_count(unfiltered, self.record_count);
225            if excluded > 0 {
226                warn!(
227                    "{excluded} blocks excluded from {desc} by the audience-mismatch predicate \
228                     -- see block_audience_mismatch_rows for per-block detail"
229                );
230                imetric!("block_audience_mismatch_excluded", "count", excluded as u64);
231            }
232        }
233
234        let (tx, rx) = tokio::sync::mpsc::channel(1);
235        let join_handle = spawn_with_context(write_partition_from_rows(
236            lake.clone(),
237            self.view_metadata.clone(),
238            self.schema.clone(),
239            self.insert_range,
240            self.get_source_data_hash(),
241            self.sort_order.clone(),
242            RetireMatch::Containment,
243            Vec::new(),
244            rx,
245            logger.clone(),
246        ));
247
248        let stream_result: Result<()> = instrument_named!(
249            async {
250                if self.record_count > 0 {
251                    let mut query = sqlx::query(&self.data_sql)
252                        .bind(self.insert_range.begin)
253                        .bind(self.insert_range.end);
254                    if let Some(default_audience) = self.default_audience.as_deref() {
255                        query = query.bind(default_audience.to_owned());
256                    }
257                    let mut rows = query.fetch(&lake.db_pool);
258                    let ctx = SessionContext::new();
259                    let mut chunk: Vec<PgRow> = Vec::new();
260                    let mut chunk_bytes = 0usize;
261                    while let Some(row) = rows.try_next().await? {
262                        chunk_bytes += estimate_row_bytes(&row);
263                        chunk.push(row);
264                        if chunk_bytes >= SOURCE_BYTES_PER_BATCH {
265                            flush_chunk(
266                                &mut chunk,
267                                &ctx,
268                                &self.schema,
269                                &self.compute_time_bounds,
270                                &tx,
271                            )
272                            .await?;
273                            chunk_bytes = 0;
274                        }
275                    }
276                    if !chunk.is_empty() {
277                        flush_chunk(
278                            &mut chunk,
279                            &ctx,
280                            &self.schema,
281                            &self.compute_time_bounds,
282                            &tx,
283                        )
284                        .await?;
285                    }
286                }
287                Ok(())
288            },
289            "sql_select_partition_source_data"
290        )
291        .await;
292
293        match stream_result {
294            Ok(()) => {
295                drop(tx);
296                join_handle.await??;
297                Ok(())
298            }
299            Err(e) => {
300                // mirror create_merged_partition's error path: send the abort through the
301                // channel before dropping it, so write_partition_from_rows sees an Err item
302                // instead of a plain closed-channel end-of-stream and does not commit a
303                // partial partition.
304                let _ = tx
305                    .send(Err(anyhow::anyhow!("metadata partition stream aborted")))
306                    .await;
307                drop(tx);
308                let _ = join_handle.await;
309                Err(e)
310            }
311        }
312    }
313}
micromegas_analytics/lakehouse/metadata_partition_spec.rs

micromegas_analytics/lakehouse/
metadata_partition_spec.rs
