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, 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::{arrow::datatypes::Schema, prelude::*};
14use futures::TryStreamExt;
15use micromegas_ingestion::data_lake_connection::DataLakeConnection;
16use micromegas_tracing::prelude::*;
17use sqlx::{Row, postgres::PgRow};
18use std::sync::Arc;
19use tokio::sync::mpsc::Sender;
20
21/// Flush threshold on the estimated byte size of the pending chunk -- bounds peak memory to one
22/// ~8 MB chunk, not one day's worth of Postgres rows. Byte-based like the Parquet writer's own
23/// 100 MB flush (`write_partition.rs`), because a row-count threshold bounds nothing when a few
24/// rows carry MB-sized properties/objects_metadata payloads. Deliberately the only flush metric.
25const SOURCE_BYTES_PER_BATCH: usize = 8 * 1024 * 1024;
26
27#[derive(Debug)]
28pub struct MetadataPartitionSpec {
29    pub view_metadata: ViewMetadata,
30    pub schema: Arc<Schema>,
31    pub insert_range: TimeRange,
32    pub record_count: i64,
33    pub data_sql: Arc<String>,
34    pub compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
35    /// The sort guarantee this partition's rows will carry, per the caller's `data_sql`'s
36    /// `ORDER BY` (e.g. `Some(["insert_time"])` for `BlocksView`). Recorded on `Partition` as-is.
37    pub sort_order: Option<Vec<String>>,
38}
39
40#[expect(clippy::too_many_arguments)]
41pub async fn fetch_metadata_partition_spec(
42    pool: &sqlx::PgPool,
43    source_count_query: &str,
44    data_sql: Arc<String>,
45    view_metadata: ViewMetadata,
46    schema: Arc<Schema>,
47    insert_range: TimeRange,
48    compute_time_bounds: Arc<dyn DataFrameTimeBounds>,
49    sort_order: Option<Vec<String>>,
50) -> Result<MetadataPartitionSpec> {
51    //todo: extract this query to allow join (instead of source_table)
52    let row = instrument_named!(
53        sqlx::query(source_count_query)
54            .bind(insert_range.begin)
55            .bind(insert_range.end)
56            .fetch_one(pool),
57        "sql_select_source_count"
58    )
59    .await
60    .with_context(|| "select count source metadata")?;
61    Ok(MetadataPartitionSpec {
62        view_metadata,
63        schema,
64        insert_range,
65        record_count: row.try_get("count").with_context(|| "reading count")?,
66        data_sql,
67        compute_time_bounds,
68        sort_order,
69    })
70}
71
72/// Estimates a row's payload size by summing its raw column value byte lengths, counting `NULL`
73/// and any non-byte-backed value as 0. This deliberately tracks the JSONB/binary columns
74/// (`properties`, `objects_metadata`, `dependencies_metadata`) that dominate blocks-view row
75/// width -- an allocator-exact footprint is not needed, only a flush-decision estimate.
76fn estimate_row_bytes(row: &PgRow) -> usize {
77    let mut total = 0usize;
78    for i in 0..row.len() {
79        if let Ok(raw) = row.try_get_raw(i)
80            && let Ok(bytes) = raw.as_bytes()
81        {
82            total += bytes.len();
83        }
84    }
85    total
86}
87
88/// Converts the accumulated chunk to a `RecordBatch`, computes its event-time bounds, and sends
89/// it as a `PartitionRowSet`. Clears `chunk` in place for reuse by the next flush.
90async fn flush_chunk(
91    chunk: &mut Vec<PgRow>,
92    ctx: &SessionContext,
93    compute_time_bounds: &Arc<dyn DataFrameTimeBounds>,
94    tx: &Sender<Result<PartitionRowSet, anyhow::Error>>,
95) -> Result<()> {
96    let record_batch =
97        rows_to_record_batch(chunk).with_context(|| "converting rows to record batch")?;
98    chunk.clear();
99    let event_time_range = compute_time_bounds
100        .get_time_bounds(
101            ctx.read_batch(record_batch.clone())
102                .with_context(|| "read_batch")?,
103        )
104        .await?;
105    tx.send(Ok(PartitionRowSet::new(event_time_range, record_batch)))
106        .await
107        .with_context(|| "sending partition row set")?;
108    Ok(())
109}
110
111#[async_trait]
112impl PartitionSpec for MetadataPartitionSpec {
113    fn is_empty(&self) -> bool {
114        self.record_count < 1
115    }
116
117    fn get_source_data_hash(&self) -> Vec<u8> {
118        self.record_count.to_le_bytes().to_vec()
119    }
120
121    async fn write(&self, lake: Arc<DataLakeConnection>, logger: Arc<dyn Logger>) -> Result<()> {
122        // Allow empty record_count - write_partition_from_rows will create
123        // an empty partition record if no data is sent through the channel
124        let desc = format!(
125            "[{}, {}] {} {}",
126            self.view_metadata.view_set_name,
127            self.view_metadata.view_instance_id,
128            self.insert_range.begin.to_rfc3339(),
129            self.insert_range.end.to_rfc3339()
130        );
131        logger.write_log_entry(format!("writing {desc}")).await?;
132
133        let (tx, rx) = tokio::sync::mpsc::channel(1);
134        let join_handle = spawn_with_context(write_partition_from_rows(
135            lake.clone(),
136            self.view_metadata.clone(),
137            self.schema.clone(),
138            self.insert_range,
139            self.get_source_data_hash(),
140            self.sort_order.clone(),
141            rx,
142            logger.clone(),
143        ));
144
145        let stream_result: Result<()> = instrument_named!(
146            async {
147                if self.record_count > 0 {
148                    let mut rows = sqlx::query(&self.data_sql)
149                        .bind(self.insert_range.begin)
150                        .bind(self.insert_range.end)
151                        .fetch(&lake.db_pool);
152                    let ctx = SessionContext::new();
153                    let mut chunk: Vec<PgRow> = Vec::new();
154                    let mut chunk_bytes = 0usize;
155                    while let Some(row) = rows.try_next().await? {
156                        chunk_bytes += estimate_row_bytes(&row);
157                        chunk.push(row);
158                        if chunk_bytes >= SOURCE_BYTES_PER_BATCH {
159                            flush_chunk(&mut chunk, &ctx, &self.compute_time_bounds, &tx).await?;
160                            chunk_bytes = 0;
161                        }
162                    }
163                    if !chunk.is_empty() {
164                        flush_chunk(&mut chunk, &ctx, &self.compute_time_bounds, &tx).await?;
165                    }
166                }
167                Ok(())
168            },
169            "sql_select_partition_source_data"
170        )
171        .await;
172
173        match stream_result {
174            Ok(()) => {
175                drop(tx);
176                join_handle.await??;
177                Ok(())
178            }
179            Err(e) => {
180                // mirror create_merged_partition's error path: send the abort through the
181                // channel before dropping it, so write_partition_from_rows sees an Err item
182                // instead of a plain closed-channel end-of-stream and does not commit a
183                // partial partition.
184                let _ = tx
185                    .send(Err(anyhow::anyhow!("metadata partition stream aborted")))
186                    .await;
187                drop(tx);
188                let _ = join_handle.await;
189                Err(e)
190            }
191        }
192    }
193}
micromegas_analytics/lakehouse/metadata_partition_spec.rs

micromegas_analytics/lakehouse/
metadata_partition_spec.rs
