Skip to main content

micromegas_analytics/lakehouse/
images_view.rs

1use super::{
2    batch_update::PartitionCreationStrategy,
3    block_partition_spec::{BlockProcessor, BlockProcessorMap},
4    blocks_view::BlocksView,
5    dataframe_time_bounds::{DataFrameTimeBounds, NamedColumnsTimeBounds},
6    image_block_processor::ImageBlockProcessor,
7    jit_partitions::{
8        BlockOrder, JitPartitionConfig, find_up_to_date_partitions,
9        generate_process_jit_partitions, write_partition_from_blocks,
10    },
11    lakehouse_context::LakehouseContext,
12    partition_cache::PartitionCache,
13    view::{PartitionSpec, View, ViewMetadata},
14    view_factory::ViewMaker,
15};
16use crate::{
17    images_table::images_table_schema,
18    metadata::find_process,
19    time::{TimeRange, datetime_to_scalar},
20};
21use anyhow::{Context, Result};
22use async_trait::async_trait;
23use chrono::{DateTime, TimeDelta, Utc};
24use datafusion::{
25    arrow::datatypes::Schema,
26    logical_expr::{Between, Expr, col},
27};
28use micromegas_ingestion::web_ingestion_service::FORMAT_TRANSIT;
29use micromegas_tracing::prelude::*;
30use std::collections::HashMap;
31use std::sync::Arc;
32use uuid::Uuid;
33
34const VIEW_SET_NAME: &str = "images";
35const SCHEMA_VERSION: u8 = 2;
36
37lazy_static::lazy_static! {
38    static ref TIME_COLUMN: Arc<String> = Arc::new(String::from("time"));
39}
40
41fn image_processors() -> Arc<BlockProcessorMap> {
42    let mut m: BlockProcessorMap = HashMap::new();
43    m.insert(
44        FORMAT_TRANSIT,
45        Arc::new(ImageBlockProcessor {}) as Arc<dyn BlockProcessor>,
46    );
47    Arc::new(m)
48}
49
50#[derive(Debug)]
51pub struct ImagesViewMaker {}
52
53impl ViewMaker for ImagesViewMaker {
54    fn make_view(&self, view_instance_id: &str) -> Result<Arc<dyn View>> {
55        Ok(Arc::new(ImagesView::new(view_instance_id)?))
56    }
57
58    fn get_schema_hash(&self) -> Vec<u8> {
59        vec![SCHEMA_VERSION]
60    }
61
62    fn get_schema(&self) -> Arc<Schema> {
63        Arc::new(images_table_schema())
64    }
65}
66
67#[derive(Debug)]
68pub struct ImagesView {
69    view_set_name: Arc<String>,
70    view_instance_id: Arc<String>,
71    process_id: sqlx::types::Uuid,
72}
73
74impl ImagesView {
75    pub fn new(view_instance_id: &str) -> Result<Self> {
76        if view_instance_id == "global" {
77            anyhow::bail!("ImagesView does not support global view access");
78        }
79        let process_id = Uuid::parse_str(view_instance_id).with_context(|| "Uuid::parse_str")?;
80        Ok(Self {
81            view_set_name: Arc::new(String::from(VIEW_SET_NAME)),
82            view_instance_id: Arc::new(view_instance_id.into()),
83            process_id,
84        })
85    }
86}
87
88#[async_trait]
89impl View for ImagesView {
90    fn get_view_set_name(&self) -> Arc<String> {
91        self.view_set_name.clone()
92    }
93
94    fn get_view_instance_id(&self) -> Arc<String> {
95        self.view_instance_id.clone()
96    }
97
98    async fn make_batch_partition_spec(
99        &self,
100        _lakehouse: Arc<LakehouseContext>,
101        _existing_partitions: Arc<PartitionCache>,
102        _insert_range: TimeRange,
103    ) -> Result<Arc<dyn PartitionSpec>> {
104        anyhow::bail!("ImagesView does not support batch partition specs")
105    }
106
107    fn get_file_schema_hash(&self) -> Vec<u8> {
108        vec![SCHEMA_VERSION]
109    }
110
111    fn get_file_schema(&self) -> Arc<Schema> {
112        Arc::new(images_table_schema())
113    }
114
115    #[span_fn]
116    async fn jit_update(
117        &self,
118        lakehouse: Arc<LakehouseContext>,
119        query_range: Option<TimeRange>,
120    ) -> Result<()> {
121        let process = Arc::new(
122            find_process(
123                &lakehouse.lake().db_pool,
124                &self.process_id,
125                &lakehouse.default_audience(),
126            )
127            .await
128            .with_context(|| "find_process")?,
129        );
130        let query_range =
131            query_range.unwrap_or_else(|| TimeRange::new(process.start_time, chrono::Utc::now()));
132        let blocks_view = BlocksView::new(lakehouse.default_audience())?;
133        let all_partitions = generate_process_jit_partitions(
134            &JitPartitionConfig::default(),
135            lakehouse.clone(),
136            &blocks_view,
137            &query_range,
138            process.clone(),
139            "image",
140        )
141        .await
142        .with_context(|| "generate_process_jit_partitions")?;
143        let view_meta = ViewMetadata {
144            view_set_name: self.get_view_set_name(),
145            view_instance_id: self.get_view_instance_id(),
146            file_schema_hash: self.get_file_schema_hash(),
147        };
148        let block_processors = image_processors();
149        let up_to_date = find_up_to_date_partitions(
150            &lakehouse.lake().db_pool,
151            view_meta.clone(),
152            BlockOrder::InsertTime,
153            &all_partitions,
154        )
155        .await
156        .with_context(|| "find_up_to_date_partitions")?;
157        for (part, up_to_date) in all_partitions.into_iter().zip(up_to_date) {
158            if !up_to_date {
159                write_partition_from_blocks(
160                    lakehouse.lake().clone(),
161                    view_meta.clone(),
162                    self.get_file_schema(),
163                    part,
164                    block_processors.clone(),
165                )
166                .await
167                .with_context(|| "write_partition_from_blocks")?;
168            }
169        }
170        Ok(())
171    }
172
173    fn make_time_filter(&self, begin: DateTime<Utc>, end: DateTime<Utc>) -> Result<Vec<Expr>> {
174        Ok(vec![Expr::Between(Between::new(
175            col("time").into(),
176            false,
177            Expr::Literal(datetime_to_scalar(begin), None).into(),
178            Expr::Literal(datetime_to_scalar(end), None).into(),
179        ))])
180    }
181
182    fn get_time_bounds(&self) -> Arc<dyn DataFrameTimeBounds> {
183        Arc::new(NamedColumnsTimeBounds::new(
184            TIME_COLUMN.clone(),
185            TIME_COLUMN.clone(),
186        ))
187    }
188
189    fn get_update_group(&self) -> Option<i32> {
190        None
191    }
192
193    fn get_max_partition_time_delta(&self, _strategy: &PartitionCreationStrategy) -> TimeDelta {
194        TimeDelta::hours(1)
195    }
196}