micromegas_analytics/lakehouse/
partition_source_data.rs

1use super::blocks_view::blocks_file_schema_hash;
2use super::lakehouse_context::LakehouseContext;
3use super::partition_cache::PartitionCache;
4use crate::dfext::{
5    string_column_accessor::string_column_by_name, typed_column::typed_column_by_name,
6};
7use crate::metadata::{ProcessMetadata, StreamMetadata};
8use crate::properties::properties_column_accessor::properties_column_by_name;
9use crate::time::TimeRange;
10use crate::{
11    dfext::typed_column::typed_column,
12    lakehouse::{blocks_view::blocks_view_schema, query::query_partitions},
13};
14use anyhow::{Context, Result};
15use async_stream::try_stream;
16use async_trait::async_trait;
17use chrono::DateTime;
18use datafusion::functions_aggregate::{count::count_all, expr_fn::sum, min_max::max};
19use datafusion::{
20    arrow::array::{
21        Array, BinaryArray, GenericListArray, Int32Array, Int64Array, StringArray,
22        TimestampNanosecondArray,
23    },
24    prelude::*,
25};
26use futures::{StreamExt, stream::BoxStream};
27use micromegas_telemetry::types::block::BlockMetadata;
28use std::fmt::Debug;
29use std::sync::Arc;
30use uuid::Uuid;
31
32/// Represents a single block of source data for a partition.
33#[derive(Debug)]
34pub struct PartitionSourceBlock {
35    pub block: BlockMetadata,
36    pub stream: Arc<StreamMetadata>,
37    pub process: Arc<ProcessMetadata>,
38    /// Wire-format identifier for the block payload — keys per-block dispatch
39    /// in `BlockPartitionSpec` and `write_partition_from_blocks`. Examples:
40    /// `"micromegas-transit"`, `"otlp/v1/logs"`, `"otlp/v1/metrics"`, `"otlp/v1/traces"`.
41    pub format: String,
42}
43
44/// A trait for providing blocks of source data for partitions.
45#[async_trait]
46pub trait PartitionBlocksSource: Sync + Send + Debug {
47    /// Returns true if there are no blocks in the source.
48    fn is_empty(&self) -> bool;
49    /// Returns the number of blocks in the source.
50    fn get_nb_blocks(&self) -> i64;
51    /// Returns the maximum payload size of the blocks in the source.
52    fn get_max_payload_size(&self) -> i64;
53    /// Returns a hash of the source data.
54    fn get_source_data_hash(&self) -> Vec<u8>;
55    /// Returns a stream of the source blocks.
56    async fn get_blocks_stream(&self) -> BoxStream<'static, Result<Arc<PartitionSourceBlock>>>;
57}
58
59/// A `PartitionBlocksSource` implementation that stores blocks in memory.
60#[derive(Debug)]
61pub struct SourceDataBlocksInMemory {
62    pub blocks: Vec<Arc<PartitionSourceBlock>>,
63    pub block_ids_hash: Vec<u8>,
64}
65
66#[async_trait]
67impl PartitionBlocksSource for SourceDataBlocksInMemory {
68    fn is_empty(&self) -> bool {
69        self.blocks.is_empty()
70    }
71
72    fn get_nb_blocks(&self) -> i64 {
73        self.blocks.len() as i64
74    }
75
76    fn get_max_payload_size(&self) -> i64 {
77        let mut max_size = self.blocks[0].block.payload_size;
78        for block in &self.blocks {
79            max_size = max_size.max(block.block.payload_size);
80        }
81        max_size
82    }
83
84    fn get_source_data_hash(&self) -> Vec<u8> {
85        self.block_ids_hash.clone()
86    }
87
88    async fn get_blocks_stream(&self) -> BoxStream<'static, Result<Arc<PartitionSourceBlock>>> {
89        let stream = futures::stream::iter(self.blocks.clone()).map(Ok);
90        stream.boxed()
91    }
92}
93
94/// A `PartitionBlocksSource` implementation that fetches blocks from a DataFusion DataFrame.
95#[derive(Debug)]
96pub struct SourceDataBlocks {
97    pub blocks_dataframe: DataFrame,
98    pub object_count: i64,
99    pub block_count: i64,
100    pub max_payload: i64,
101}
102
103#[async_trait]
104impl PartitionBlocksSource for SourceDataBlocks {
105    fn is_empty(&self) -> bool {
106        self.object_count == 0
107    }
108
109    fn get_nb_blocks(&self) -> i64 {
110        self.block_count
111    }
112
113    fn get_max_payload_size(&self) -> i64 {
114        self.max_payload
115    }
116
117    fn get_source_data_hash(&self) -> Vec<u8> {
118        self.object_count.to_le_bytes().to_vec()
119    }
120
121    async fn get_blocks_stream(&self) -> BoxStream<'static, Result<Arc<PartitionSourceBlock>>> {
122        let df = self.blocks_dataframe.clone();
123        Box::pin(try_stream! {
124            let mut stream = df.execute_stream().await?;
125            while let Some(res) = stream.next().await {
126                let b = res.with_context(|| "fetching blocks query results")?;
127                let block_id_column = string_column_by_name(&b, "block_id")?;
128                let stream_id_column = string_column_by_name(&b, "stream_id")?;
129                let process_id_column = string_column_by_name(&b, "process_id")?;
130                let begin_time_column: &TimestampNanosecondArray = typed_column_by_name(&b, "begin_time")?;
131                let begin_ticks_column: &Int64Array = typed_column_by_name(&b, "begin_ticks")?;
132                let end_time_column: &TimestampNanosecondArray = typed_column_by_name(&b, "end_time")?;
133                let end_ticks_column: &Int64Array = typed_column_by_name(&b, "end_ticks")?;
134                let nb_objects_column: &Int32Array = typed_column_by_name(&b, "nb_objects")?;
135                let object_offset_column: &Int64Array = typed_column_by_name(&b, "object_offset")?;
136                let payload_size_column: &Int64Array = typed_column_by_name(&b, "payload_size")?;
137                let block_insert_time_column: &TimestampNanosecondArray =
138                    typed_column_by_name(&b, "insert_time")?;
139                let dependencies_metadata_column: &BinaryArray =
140                    typed_column_by_name(&b, "streams.dependencies_metadata")?;
141                let objects_metadata_column: &BinaryArray =
142                    typed_column_by_name(&b, "streams.objects_metadata")?;
143                let stream_tags_column: &GenericListArray<i32> = typed_column_by_name(&b, "streams.tags")?;
144                let stream_properties_accessor = properties_column_by_name(&b, "streams.properties")?;
145                let stream_format_column = string_column_by_name(&b, "streams.format")?;
146
147                let process_start_time_column: &TimestampNanosecondArray =
148                    typed_column_by_name(&b, "processes.start_time")?;
149                let process_start_ticks_column: &Int64Array =
150                    typed_column_by_name(&b, "processes.start_ticks")?;
151                let process_tsc_freq_column: &Int64Array =
152                    typed_column_by_name(&b, "processes.tsc_frequency")?;
153                let process_exe_column = string_column_by_name(&b, "processes.exe")?;
154                let process_username_column = string_column_by_name(&b, "processes.username")?;
155                let process_realname_column = string_column_by_name(&b, "processes.realname")?;
156                let process_computer_column = string_column_by_name(&b, "processes.computer")?;
157                let process_distro_column = string_column_by_name(&b, "processes.distro")?;
158                let process_cpu_column = string_column_by_name(&b, "processes.cpu_brand")?;
159                let process_parent_column = string_column_by_name(&b, "processes.parent_process_id")?;
160                let process_properties_accessor = properties_column_by_name(&b, "processes.properties")?;
161                let audience_column = string_column_by_name(&b, "audience")?;
162                for ir in 0..b.num_rows() {
163                    let block_insert_time = block_insert_time_column.value(ir);
164                    let stream_id = Uuid::parse_str(stream_id_column.value(ir)?)?;
165                    let process_id = Uuid::parse_str(process_id_column.value(ir)?)?;
166                    let block = BlockMetadata {
167                        block_id: Uuid::parse_str(block_id_column.value(ir)?)?,
168                        stream_id,
169                        process_id,
170                        begin_time: DateTime::from_timestamp_nanos(begin_time_column.value(ir)),
171                        end_time: DateTime::from_timestamp_nanos(end_time_column.value(ir)),
172                        begin_ticks: begin_ticks_column.value(ir),
173                        end_ticks: end_ticks_column.value(ir),
174                        nb_objects: nb_objects_column.value(ir),
175                        payload_size: payload_size_column.value(ir),
176                        object_offset: object_offset_column.value(ir),
177                        insert_time: DateTime::from_timestamp_nanos(block_insert_time),
178                    };
179
180                    let dependencies_metadata = dependencies_metadata_column.value(ir);
181                    let objects_metadata = objects_metadata_column.value(ir);
182                    let stream_tags = stream_tags_column
183                        .value(ir)
184                        .as_any()
185                        .downcast_ref::<StringArray>()
186                        .with_context(|| "casting stream_tags")?
187                        .iter()
188                        .map(|item| String::from(item.unwrap_or_default()))
189                        .collect();
190                    let stream_properties_jsonb = stream_properties_accessor.jsonb_value(ir)?;
191                    let stream = StreamMetadata {
192                        process_id,
193                        stream_id,
194                        dependencies_metadata: ciborium::from_reader(dependencies_metadata)
195                            .with_context(|| "decoding dependencies_metadata")?,
196                        objects_metadata: ciborium::from_reader(objects_metadata)
197                            .with_context(|| "decoding objects_metadata")?,
198                        tags: stream_tags,
199                        properties: Arc::new(stream_properties_jsonb),
200                    };
201                    let process_properties_jsonb = process_properties_accessor.jsonb_value(ir)?;
202                    let parent_value = process_parent_column.value(ir)?;
203                    let parent_process_id = if parent_value.is_empty() {
204                        None
205                    } else {
206                        Some(Uuid::parse_str(parent_value).with_context(|| "parsing parent process_id")?)
207                    };
208
209                    let process = ProcessMetadata {
210                        process_id,
211                        exe: process_exe_column.value(ir)?.into(),
212                        username: process_username_column.value(ir)?.into(),
213                        realname: process_realname_column.value(ir)?.into(),
214                        computer: process_computer_column.value(ir)?.into(),
215                        distro: process_distro_column.value(ir)?.into(),
216                        cpu_brand: process_cpu_column.value(ir)?.into(),
217                        tsc_frequency: process_tsc_freq_column.value(ir),
218                        start_time: DateTime::from_timestamp_nanos(process_start_time_column.value(ir)),
219                        start_ticks: process_start_ticks_column.value(ir),
220                        parent_process_id,
221                        properties: Arc::new(process_properties_jsonb),
222                        // The source column is non-nullable, so no null check is needed here.
223                        // That rests entirely on `BlocksView`'s `COALESCE`: the accessor does
224                        // not null-check, and the dictionary accessor would return a wrong
225                        // non-null string on a null slot rather than failing.
226                        audience: Arc::from(audience_column.value(ir)?),
227                    };
228                    let format = stream_format_column.value(ir)?.to_string();
229                    yield Arc::new(PartitionSourceBlock {
230                        block,
231                        stream: Arc::new(stream),
232                        process: Arc::new(process),
233                        format,
234                    });
235                }
236            }
237
238        })
239    }
240}
241
242/// Converts a hash (expected to be an i64 as bytes) to an object count.
243pub fn hash_to_object_count(hash: &[u8]) -> Result<i64> {
244    Ok(i64::from_le_bytes(
245        hash.try_into().with_context(|| "hash_to_object_count")?,
246    ))
247}
248
249/// Fetches partition source data from the data lake.
250pub async fn fetch_partition_source_data(
251    lakehouse: Arc<LakehouseContext>,
252    existing_partitions: Arc<PartitionCache>,
253    insert_range: TimeRange,
254    source_stream_tag: &str,
255) -> Result<SourceDataBlocks> {
256    let begin_rfc = insert_range.begin.to_rfc3339();
257    let end_rfc = insert_range.end.to_rfc3339();
258    let sql = format!(
259        r#"
260          SELECT block_id, stream_id, process_id, begin_time, begin_ticks, end_time, end_ticks, nb_objects,
261              object_offset, payload_size, insert_time,
262              "streams.dependencies_metadata", "streams.objects_metadata", "streams.tags", "streams.properties", "streams.format",
263              "processes.start_time", "processes.start_ticks", "processes.tsc_frequency", "processes.exe",
264              "processes.username", "processes.realname", "processes.computer", "processes.distro", "processes.cpu_brand",
265              "processes.parent_process_id", "processes.properties", audience
266          FROM source
267          WHERE array_has( "streams.tags", '{source_stream_tag}' )
268          AND insert_time >= '{begin_rfc}'
269          AND insert_time < '{end_rfc}'
270          ;"#
271    );
272    let block_partitions = existing_partitions
273        .filter("blocks", "global", &blocks_file_schema_hash(), insert_range)
274        .partitions;
275    let reader_factory = lakehouse.reader_factory().clone();
276    let df = query_partitions(
277        lakehouse.runtime().clone(),
278        reader_factory,
279        lakehouse.lake().blob_storage.inner(),
280        Arc::new(blocks_view_schema()),
281        Arc::new(block_partitions),
282        &sql,
283    )
284    .await
285    .with_context(|| "blocks query")?;
286    let blocks_stats_df = df.clone().aggregate(
287        vec![],
288        vec![
289            sum(col("nb_objects")),
290            count_all(),
291            max(col("payload_size")),
292        ],
293    )?;
294    let blocks_stats_rbs = blocks_stats_df.collect().await?;
295    if blocks_stats_rbs.len() != 1 {
296        anyhow::bail!("nb_objects_rbs has size {}", blocks_stats_rbs.len());
297    }
298    if blocks_stats_rbs[0].num_rows() != 1 {
299        anyhow::bail!(
300            "nb_objects_rbs[0] has size {}",
301            blocks_stats_rbs[0].num_rows()
302        );
303    }
304    let sub_nb_objects_column: &Int64Array = typed_column(&blocks_stats_rbs[0], 0)?;
305    let object_count = sub_nb_objects_column.value(0);
306    let block_count_column: &Int64Array = typed_column(&blocks_stats_rbs[0], 1)?;
307    let block_count = block_count_column.value(0);
308    let max_payload_size_column: &Int64Array = typed_column(&blocks_stats_rbs[0], 2)?;
309    let max_payload = max_payload_size_column.value(0);
310
311    Ok(SourceDataBlocks {
312        blocks_dataframe: df,
313        object_count,
314        block_count,
315        max_payload,
316    })
317}
micromegas_analytics/lakehouse/partition_source_data.rs

micromegas_analytics/lakehouse/
partition_source_data.rs
