micromegas_analytics/lakehouse/
batch_partition_merger.rs

1use super::{
2    lakehouse_context::LakehouseContext,
3    merge::{MergeQueryResult, PartitionMerger},
4    partition::Partition,
5    partition_cache::PartitionCache,
6    session_configurator::SessionConfigurator,
7    view_factory::ViewFactory,
8};
9use crate::{
10    lakehouse::{
11        partitioned_table_provider::PartitionedTableProvider, query::make_session_context,
12    },
13    time::{TimeRange, datetime_to_scalar},
14};
15use anyhow::Result;
16use async_trait::async_trait;
17use chrono::{DateTime, TimeDelta, Utc};
18use datafusion::{
19    arrow::datatypes::Schema, error::DataFusionError,
20    physical_plan::stream::RecordBatchReceiverStreamBuilder, sql::TableReference,
21};
22use futures::TryStreamExt;
23use futures::{StreamExt, stream};
24use micromegas_tracing::prelude::*;
25use std::sync::Arc;
26
27/// Statistics about a set of partitions.
28struct PartitionStats {
29    pub num_rows: i64,
30    pub min_event_time: DateTime<Utc>,
31    pub max_event_time: DateTime<Utc>,
32}
33
34fn compute_partition_stats(partitions: &[Partition]) -> Result<PartitionStats> {
35    // Filter out empty partitions before computing stats
36    let non_empty: Vec<_> = partitions.iter().filter(|p| !p.is_empty()).collect();
37
38    if non_empty.is_empty() {
39        anyhow::bail!(
40            "compute_partition_stats given only empty partitions (should be filtered at caller)"
41        );
42    }
43
44    let first = non_empty.first().unwrap();
45    let first_event_range = first
46        .event_time_range
47        .ok_or_else(|| anyhow::anyhow!("non-empty partition has no event_time_range"))?;
48
49    let state = PartitionStats {
50        num_rows: first.num_rows,
51        min_event_time: first_event_range.begin,
52        max_event_time: first_event_range.end,
53    };
54
55    non_empty
56        .iter()
57        .skip(1)
58        .try_fold(state, |state, part| -> Result<PartitionStats> {
59            let event_range = part
60                .event_time_range
61                .ok_or_else(|| anyhow::anyhow!("non-empty partition has no event_time_range"))?;
62            Ok(PartitionStats {
63                num_rows: state.num_rows + part.num_rows,
64                min_event_time: state.min_event_time.min(event_range.begin),
65                max_event_time: state.max_event_time.max(event_range.end),
66            })
67        })
68}
69
70/// Merges multiple partitions by splitting the work in batches to use less memory.
71/// The batches are based on event times.
72#[derive(Debug)]
73pub struct BatchPartitionMerger {
74    /// file_schema: arrow schema of the parquet files
75    file_schema: Arc<Schema>,
76    /// view_factory: allows joins in merge query
77    view_factory: Arc<ViewFactory>,
78    /// session_configurator: allows custom tables in merge query
79    session_configurator: Arc<dyn SessionConfigurator>,
80    /// merge_batch_query: merge query with begin & end placeholders
81    merge_batch_query: String,
82    /// batch size to aim for
83    approx_nb_rows_per_batch: i64,
84}
85
86impl BatchPartitionMerger {
87    pub fn new(
88        file_schema: Arc<Schema>,
89        view_factory: Arc<ViewFactory>,
90        session_configurator: Arc<dyn SessionConfigurator>,
91        merge_batch_query: String,
92        approx_nb_rows_per_batch: i64,
93    ) -> Self {
94        Self {
95            file_schema,
96            view_factory,
97            session_configurator,
98            merge_batch_query,
99            approx_nb_rows_per_batch,
100        }
101    }
102}
103
104#[async_trait]
105impl PartitionMerger for BatchPartitionMerger {
106    #[span_fn]
107    async fn execute_merge_query(
108        &self,
109        lakehouse: Arc<LakehouseContext>,
110        partitions_to_merge: Arc<Vec<Partition>>,
111        partitions_all_views: Arc<PartitionCache>,
112        insert_range: TimeRange,
113    ) -> Result<MergeQueryResult> {
114        info!("execute_merge_query");
115
116        // If all partitions are empty, return empty stream immediately
117        if partitions_to_merge.iter().all(|p| p.is_empty()) {
118            debug!("all partitions are empty, returning empty stream");
119            let builder = RecordBatchReceiverStreamBuilder::new(self.file_schema.clone(), 1);
120            return Ok(MergeQueryResult {
121                stream: builder.build(),
122                ordering_honored: true,
123            });
124        }
125
126        let stats = compute_partition_stats(partitions_to_merge.as_ref())?;
127        let nb_batches = ((stats.num_rows / self.approx_nb_rows_per_batch) + 1) as i32;
128        let batch_time_delta = ((stats.max_event_time - stats.min_event_time) / nb_batches)
129            + TimeDelta::nanoseconds(1);
130
131        let file_schema = self.file_schema.clone();
132        let reader_factory = lakehouse.reader_factory().clone();
133        let ctx = make_session_context(
134            lakehouse.clone(),
135            partitions_all_views,
136            Some(insert_range),
137            self.view_factory.clone(),
138            self.session_configurator.clone(),
139            true,
140        )
141        .await?;
142        let src_table =
143            PartitionedTableProvider::new(file_schema, reader_factory, partitions_to_merge);
144        ctx.register_table(
145            TableReference::Bare {
146                table: "source".into(),
147            },
148            Arc::new(src_table),
149        )?;
150        let df_template = ctx.sql(&self.merge_batch_query).await.map_err(|e| {
151            DataFusionError::Execution(format!("building template for merge query: {e:?}"))
152        })?;
153
154        let mut builder = RecordBatchReceiverStreamBuilder::new(self.file_schema.clone(), 10);
155        let sender = builder.tx();
156        builder.spawn(async move {
157            let mut streams_stream = stream::iter(0..nb_batches)
158                .map(|i| {
159                    let begin = stats.min_event_time + (batch_time_delta * i);
160                    let end = begin + batch_time_delta;
161                    debug!("merging batch {begin} {end}");
162                    df_template
163                        .clone()
164                        .with_param_values(vec![
165                            ("begin", datetime_to_scalar(begin)),
166                            ("end", datetime_to_scalar(end)),
167                        ])
168                        .map(|df| async {
169                            spawn_with_context(df.execute_stream())
170                                .await
171                                .map_err(|e| DataFusionError::External(e.into()))
172                        })
173                })
174                .try_buffered(2);
175
176            while let Some(stream_res) = streams_stream.next().await {
177                let mut merge_stream = stream_res??;
178                let sender = sender.clone();
179                while let Some(rb_res) = merge_stream.next().await {
180                    if let Err(e) = sender.send(rb_res).await {
181                        error!("sending record batch: {e:?}");
182                    }
183                }
184            }
185            Ok(())
186        });
187        debug!("building merge stream");
188        Ok(MergeQueryResult {
189            stream: builder.build(),
190            ordering_honored: true,
191        })
192    }
193}
micromegas_analytics/lakehouse/batch_partition_merger.rs

micromegas_analytics/lakehouse/
batch_partition_merger.rs
