3929938005

Build Type

push

github

Committed by GitHub

Commit Message

Merge #713

Run Details

84930 of 96741 relevant lines covered (87.79%)

79640.1 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

86.99

/operators/src/engine/execution_context.rs

use super::query::QueryAbortRegistration;
use super::{
    CreateSpan, InitializedPlotOperator, InitializedRasterOperator, InitializedVectorOperator,
    MockQueryContext,
};
use crate::engine::{
    ChunkByteSize, RasterResultDescriptor, ResultDescriptor, VectorResultDescriptor,
};
use crate::error::Error;
use crate::mock::MockDatasetDataSourceLoadingInfo;
use crate::source::{GdalLoadingInfo, OgrSourceDataset};
use crate::util::{create_rayon_thread_pool, Result};
use async_trait::async_trait;
use geoengine_datatypes::dataset::DataId;
use geoengine_datatypes::primitives::{RasterQueryRectangle, VectorQueryRectangle};
use geoengine_datatypes::raster::TilingSpecification;
use geoengine_datatypes::util::test::TestDefault;
use rayon::ThreadPool;
use serde::{Deserialize, Serialize};
use std::any::Any;
use std::collections::HashMap;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::path::PathBuf;
use std::sync::Arc;

/// A context that provides certain utility access during operator initialization
#[async_trait::async_trait]
pub trait ExecutionContext: Send
    + Sync
    + MetaDataProvider<MockDatasetDataSourceLoadingInfo, VectorResultDescriptor, VectorQueryRectangle>
    + MetaDataProvider<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>
    + MetaDataProvider<GdalLoadingInfo, RasterResultDescriptor, RasterQueryRectangle>
{
    fn thread_pool(&self) -> &Arc<ThreadPool>;
    fn tiling_specification(&self) -> TilingSpecification;

    fn wrap_initialized_raster_operator(
        &self,
        op: Box<dyn InitializedRasterOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedRasterOperator>;

    fn wrap_initialized_vector_operator(
        &self,
        op: Box<dyn InitializedVectorOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedVectorOperator>;

    fn wrap_initialized_plot_operator(
        &self,
        op: Box<dyn InitializedPlotOperator>,
        span: CreateSpan,
    ) -> Box<dyn InitializedPlotOperator>;

    async fn read_ml_model(&self, path: PathBuf) -> Result<String>;

    async fn write_ml_model(&mut self, path: PathBuf, ml_model_str: String) -> Result<()>;
}

#[async_trait]
pub trait MetaDataProvider<L, R, Q>
where
    R: ResultDescriptor,
{
    async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>>;
}

#[async_trait]
pub trait MetaData<L, R, Q>: Debug + Send + Sync
where
    R: ResultDescriptor,
{
    async fn loading_info(&self, query: Q) -> Result<L>;
    async fn result_descriptor(&self) -> Result<R>;

    fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>>;
}

impl<L, R, Q> Clone for Box<dyn MetaData<L, R, Q>>
where
    R: ResultDescriptor,
{
    fn clone(&self) -> Box<dyn MetaData<L, R, Q>> {
        self.box_clone()
    }
}

pub struct MockExecutionContext {
    pub thread_pool: Arc<ThreadPool>,
    pub meta_data: HashMap<DataId, Box<dyn Any + Send + Sync>>,
    pub tiling_specification: TilingSpecification,
    pub ml_models: HashMap<PathBuf, String>,
}

impl TestDefault for MockExecutionContext {
    fn test_default() -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(0),
            meta_data: HashMap::default(),
            tiling_specification: TilingSpecification::test_default(),
            ml_models: HashMap::default(),
        }
    }
}

impl MockExecutionContext {
    pub fn new_with_tiling_spec(tiling_specification: TilingSpecification) -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(0),
            meta_data: HashMap::default(),
            tiling_specification,
            ml_models: HashMap::default(),
        }
    }

    pub fn new_with_tiling_spec_and_thread_count(
        tiling_specification: TilingSpecification,
        num_threads: usize,
    ) -> Self {
        Self {
            thread_pool: create_rayon_thread_pool(num_threads),
            meta_data: HashMap::default(),
            tiling_specification,
            ml_models: HashMap::default(),
        }
    }

    pub fn add_meta_data<L, R, Q>(&mut self, data: DataId, meta_data: Box<dyn MetaData<L, R, Q>>)
    where
        L: Send + Sync + 'static,
        R: Send + Sync + 'static + ResultDescriptor,
        Q: Send + Sync + 'static,
    {
        self.meta_data
            .insert(data, Box::new(meta_data) as Box<dyn Any + Send + Sync>);
    }

    pub fn mock_query_context(&self, chunk_byte_size: ChunkByteSize) -> MockQueryContext {
        let (abort_registration, abort_trigger) = QueryAbortRegistration::new();
        MockQueryContext {
            chunk_byte_size,
            thread_pool: self.thread_pool.clone(),
            extensions: Default::default(),
            abort_registration,
            abort_trigger: Some(abort_trigger),
        }
    }

    pub fn initialize_ml_model(&mut self, model_path: PathBuf) -> Result<()> {
        let model = std::fs::read_to_string(&model_path)?;

        self.ml_models.insert(model_path, model);

        Ok(())
    }
}

#[async_trait::async_trait]
impl ExecutionContext for MockExecutionContext {
    fn thread_pool(&self) -> &Arc<ThreadPool> {
        &self.thread_pool
    }

    fn tiling_specification(&self) -> TilingSpecification {
        self.tiling_specification
    }

    fn wrap_initialized_raster_operator(
        &self,
        op: Box<dyn InitializedRasterOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedRasterOperator> {
        op
    }

    fn wrap_initialized_vector_operator(
        &self,
        op: Box<dyn InitializedVectorOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedVectorOperator> {
        op
    }

    fn wrap_initialized_plot_operator(
        &self,
        op: Box<dyn InitializedPlotOperator>,
        _span: CreateSpan,
    ) -> Box<dyn InitializedPlotOperator> {
        op
    }

    async fn read_ml_model(&self, path: PathBuf) -> Result<String> {
        let res = self
            .ml_models
            .get(&path)
            .ok_or(Error::MachineLearningModelNotFound)?
            .clone();

        Ok(res)
    }

    async fn write_ml_model(&mut self, path: PathBuf, ml_model_str: String) -> Result<()> {
        self.ml_models.insert(path, ml_model_str);

        Ok(())
    }
}

#[async_trait]
impl<L, R, Q> MetaDataProvider<L, R, Q> for MockExecutionContext
where
    L: 'static,
    R: 'static + ResultDescriptor,
    Q: 'static,
{
    async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>> {
        let meta_data = self
            .meta_data
            .get(id)
            .ok_or(Error::UnknownDataId)?
            .downcast_ref::<Box<dyn MetaData<L, R, Q>>>()
            .ok_or(Error::InvalidMetaDataType)?;

        Ok(meta_data.clone())
    }
}

#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StaticMetaData<L, R, Q>
where
    L: Debug + Clone + Send + Sync + 'static,
    R: Debug + Send + Sync + 'static + ResultDescriptor,
    Q: Debug + Clone + Send + Sync + 'static,
{
    pub loading_info: L,
    pub result_descriptor: R,
    #[serde(skip)]
    pub phantom: PhantomData<Q>,
}

#[async_trait]
impl<L, R, Q> MetaData<L, R, Q> for StaticMetaData<L, R, Q>
where
    L: Debug + Clone + Send + Sync + 'static,
    R: Debug + Send + Sync + 'static + ResultDescriptor,
    Q: Debug + Clone + Send + Sync + 'static,
{
    async fn loading_info(&self, _query: Q) -> Result<L> {
        Ok(self.loading_info.clone())
    }

    async fn result_descriptor(&self) -> Result<R> {
        Ok(self.result_descriptor.clone())
    }

    fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>> {
        Box::new(self.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use geoengine_datatypes::collections::VectorDataType;
    use geoengine_datatypes::spatial_reference::SpatialReferenceOption;

    #[tokio::test]
    async fn test() {
        let info = StaticMetaData {
            loading_info: 1_i32,
            result_descriptor: VectorResultDescriptor {
                data_type: VectorDataType::Data,
                spatial_reference: SpatialReferenceOption::Unreferenced,
                columns: Default::default(),
                time: None,
                bbox: None,
            },
            phantom: Default::default(),
        };

        let info: Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>> =
            Box::new(info);

        let info2: Box<dyn Any + Send + Sync> = Box::new(info);

        let info3 = info2
            .downcast_ref::<Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>>>()
            .unwrap();

        assert_eq!(
            info3.result_descriptor().await.unwrap(),
            VectorResultDescriptor {
                data_type: VectorDataType::Data,
                spatial_reference: SpatialReferenceOption::Unreferenced,
                columns: Default::default(),
                time: None,
                bbox: None,
            }
        );
    }
}

1	use super::query::QueryAbortRegistration;
2	use super::{
3	CreateSpan, InitializedPlotOperator, InitializedRasterOperator, InitializedVectorOperator,
4	MockQueryContext,
5	};
6	use crate::engine::{
7	ChunkByteSize, RasterResultDescriptor, ResultDescriptor, VectorResultDescriptor,
8	};
9	use crate::error::Error;
10	use crate::mock::MockDatasetDataSourceLoadingInfo;
11	use crate::source::{GdalLoadingInfo, OgrSourceDataset};
12	use crate::util::{create_rayon_thread_pool, Result};
13	use async_trait::async_trait;
14	use geoengine_datatypes::dataset::DataId;
15	use geoengine_datatypes::primitives::{RasterQueryRectangle, VectorQueryRectangle};
16	use geoengine_datatypes::raster::TilingSpecification;
17	use geoengine_datatypes::util::test::TestDefault;
18	use rayon::ThreadPool;
19	use serde::{Deserialize, Serialize};
20	use std::any::Any;
21	use std::collections::HashMap;
22	use std::fmt::Debug;
23	use std::marker::PhantomData;
24	use std::path::PathBuf;
25	use std::sync::Arc;
26
27	/// A context that provides certain utility access during operator initialization
28	#[async_trait::async_trait]
29	pub trait ExecutionContext: Send
30	+ Sync
31	+ MetaDataProvider<MockDatasetDataSourceLoadingInfo, VectorResultDescriptor, VectorQueryRectangle>
32	+ MetaDataProvider<OgrSourceDataset, VectorResultDescriptor, VectorQueryRectangle>
33	+ MetaDataProvider<GdalLoadingInfo, RasterResultDescriptor, RasterQueryRectangle>
34	{
35	fn thread_pool(&self) -> &Arc<ThreadPool>;
36	fn tiling_specification(&self) -> TilingSpecification;
37
38	fn wrap_initialized_raster_operator(
39	&self,
40	op: Box<dyn InitializedRasterOperator>,
41	span: CreateSpan,
42	) -> Box<dyn InitializedRasterOperator>;
43
44	fn wrap_initialized_vector_operator(
45	&self,
46	op: Box<dyn InitializedVectorOperator>,
47	span: CreateSpan,
48	) -> Box<dyn InitializedVectorOperator>;
49
50	fn wrap_initialized_plot_operator(
51	&self,
52	op: Box<dyn InitializedPlotOperator>,
53	span: CreateSpan,
54	) -> Box<dyn InitializedPlotOperator>;
55
56	async fn read_ml_model(&self, path: PathBuf) -> Result<String>;
57
58	async fn write_ml_model(&mut self, path: PathBuf, ml_model_str: String) -> Result<()>;
59	}
60
61	#[async_trait]
62	pub trait MetaDataProvider<L, R, Q>
63	where
64	R: ResultDescriptor,
65	{
66	async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>>;
67	}
68
69	#[async_trait]
70	pub trait MetaData<L, R, Q>: Debug + Send + Sync
71	where
72	R: ResultDescriptor,
73	{
74	async fn loading_info(&self, query: Q) -> Result<L>;
75	async fn result_descriptor(&self) -> Result<R>;
76
77	fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>>;
78	}
79
80	impl<L, R, Q> Clone for Box<dyn MetaData<L, R, Q>>
81	where
82	R: ResultDescriptor,
83	{
84	fn clone(&self) -> Box<dyn MetaData<L, R, Q>> {	125✔
85	self.box_clone()	125✔
86	}	125✔
87	}
88
89	pub struct MockExecutionContext {
90	pub thread_pool: Arc<ThreadPool>,
91	pub meta_data: HashMap<DataId, Box<dyn Any + Send + Sync>>,
92	pub tiling_specification: TilingSpecification,
93	pub ml_models: HashMap<PathBuf, String>,
94	}
95
96	impl TestDefault for MockExecutionContext {
97	fn test_default() -> Self {	116✔
98	Self {	116✔
99	thread_pool: create_rayon_thread_pool(0),	116✔
100	meta_data: HashMap::default(),	116✔
101	tiling_specification: TilingSpecification::test_default(),	116✔
102	ml_models: HashMap::default(),	116✔
103	}	116✔
104	}	116✔
105	}
106
107	impl MockExecutionContext {
108	pub fn new_with_tiling_spec(tiling_specification: TilingSpecification) -> Self {	119✔
109	Self {	119✔
110	thread_pool: create_rayon_thread_pool(0),	119✔
111	meta_data: HashMap::default(),	119✔
112	tiling_specification,	119✔
113	ml_models: HashMap::default(),	119✔
114	}	119✔
115	}	119✔
116
117	pub fn new_with_tiling_spec_and_thread_count(	×
118	tiling_specification: TilingSpecification,	×
119	num_threads: usize,	×
120	) -> Self {	×
121	Self {	×
122	thread_pool: create_rayon_thread_pool(num_threads),	×
123	meta_data: HashMap::default(),	×
124	tiling_specification,	×
125	ml_models: HashMap::default(),	×
126	}	×
127	}	×
128
129	pub fn add_meta_data<L, R, Q>(&mut self, data: DataId, meta_data: Box<dyn MetaData<L, R, Q>>)	46✔
130	where	46✔
131	L: Send + Sync + 'static,	46✔
132	R: Send + Sync + 'static + ResultDescriptor,	46✔
133	Q: Send + Sync + 'static,	46✔
134	{	46✔
135	self.meta_data	46✔
136	.insert(data, Box::new(meta_data) as Box<dyn Any + Send + Sync>);	46✔
137	}	46✔
138
139	pub fn mock_query_context(&self, chunk_byte_size: ChunkByteSize) -> MockQueryContext {	3✔
140	let (abort_registration, abort_trigger) = QueryAbortRegistration::new();	3✔
141	MockQueryContext {	3✔
142	chunk_byte_size,	3✔
143	thread_pool: self.thread_pool.clone(),	3✔
144	extensions: Default::default(),	3✔
145	abort_registration,	3✔
146	abort_trigger: Some(abort_trigger),	3✔
147	}	3✔
148	}	3✔
149
150	pub fn initialize_ml_model(&mut self, model_path: PathBuf) -> Result<()> {	2✔
151	let model = std::fs::read_to_string(&model_path)?;	2✔
152
153	self.ml_models.insert(model_path, model);	2✔
154		2✔
155	Ok(())	2✔
156	}	2✔
157	}
158
159	#[async_trait::async_trait]
160	impl ExecutionContext for MockExecutionContext {
161	fn thread_pool(&self) -> &Arc<ThreadPool> {	×
162	&self.thread_pool	×
163	}	×
164
165	fn tiling_specification(&self) -> TilingSpecification {	195✔
166	self.tiling_specification	195✔
167	}	195✔
168
169	fn wrap_initialized_raster_operator(	238✔
170	&self,	238✔
171	op: Box<dyn InitializedRasterOperator>,	238✔
172	_span: CreateSpan,	238✔
173	) -> Box<dyn InitializedRasterOperator> {	238✔
174	op	238✔
175	}	238✔
176
177	fn wrap_initialized_vector_operator(	146✔
178	&self,	146✔
179	op: Box<dyn InitializedVectorOperator>,	146✔
180	_span: CreateSpan,	146✔
181	) -> Box<dyn InitializedVectorOperator> {	146✔
182	op	146✔
183	}	146✔
184
185	fn wrap_initialized_plot_operator(	51✔
186	&self,	51✔
187	op: Box<dyn InitializedPlotOperator>,	51✔
188	_span: CreateSpan,	51✔
189	) -> Box<dyn InitializedPlotOperator> {	51✔
190	op	51✔
191	}	51✔
192
193	async fn read_ml_model(&self, path: PathBuf) -> Result<String> {	2✔
194	let res = self	2✔
195	.ml_models	2✔
196	.get(&path)	2✔
197	.ok_or(Error::MachineLearningModelNotFound)?	2✔
198	.clone();	2✔
199		2✔
200	Ok(res)	2✔
201	}	4✔
202
203	async fn write_ml_model(&mut self, path: PathBuf, ml_model_str: String) -> Result<()> {	×
204	self.ml_models.insert(path, ml_model_str);	×
205		×
206	Ok(())	×
207	}	×
208	}
209
210	#[async_trait]
211	impl<L, R, Q> MetaDataProvider<L, R, Q> for MockExecutionContext
212	where
213	L: 'static,
214	R: 'static + ResultDescriptor,
215	Q: 'static,
216	{
217	async fn meta_data(&self, id: &DataId) -> Result<Box<dyn MetaData<L, R, Q>>> {	48✔
218	let meta_data = self	48✔
219	.meta_data	48✔
220	.get(id)	48✔
221	.ok_or(Error::UnknownDataId)?	48✔
222	.downcast_ref::<Box<dyn MetaData<L, R, Q>>>()	48✔
223	.ok_or(Error::InvalidMetaDataType)?;	48✔
224
225	Ok(meta_data.clone())	48✔
226	}	96✔
227	}
228
229	#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]	67✔
230	#[serde(rename_all = "camelCase")]
231	pub struct StaticMetaData<L, R, Q>
232	where
233	L: Debug + Clone + Send + Sync + 'static,
234	R: Debug + Send + Sync + 'static + ResultDescriptor,
235	Q: Debug + Clone + Send + Sync + 'static,
236	{
237	pub loading_info: L,
238	pub result_descriptor: R,
239	#[serde(skip)]
240	pub phantom: PhantomData<Q>,
241	}
242
243	#[async_trait]
244	impl<L, R, Q> MetaData<L, R, Q> for StaticMetaData<L, R, Q>
245	where
246	L: Debug + Clone + Send + Sync + 'static,
247	R: Debug + Send + Sync + 'static + ResultDescriptor,
248	Q: Debug + Clone + Send + Sync + 'static,
249	{
250	async fn loading_info(&self, _query: Q) -> Result<L> {	47✔
251	Ok(self.loading_info.clone())	47✔
252	}	47✔
253
254	async fn result_descriptor(&self) -> Result<R> {	32✔
255	Ok(self.result_descriptor.clone())	32✔
256	}	32✔
257
258	fn box_clone(&self) -> Box<dyn MetaData<L, R, Q>> {	47✔
259	Box::new(self.clone())	47✔
260	}	47✔
261	}
262
263	#[cfg(test)]
264	mod tests {
265	use super::*;
266	use geoengine_datatypes::collections::VectorDataType;
267	use geoengine_datatypes::spatial_reference::SpatialReferenceOption;
268
269	#[tokio::test]	1✔
270	async fn test() {	1✔
271	let info = StaticMetaData {	1✔
272	loading_info: 1_i32,	1✔
273	result_descriptor: VectorResultDescriptor {	1✔
274	data_type: VectorDataType::Data,	1✔
275	spatial_reference: SpatialReferenceOption::Unreferenced,	1✔
276	columns: Default::default(),	1✔
277	time: None,	1✔
278	bbox: None,	1✔
279	},	1✔
280	phantom: Default::default(),	1✔
281	};	1✔
282		1✔
283	let info: Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>> =	1✔
284	Box::new(info);	1✔
285		1✔
286	let info2: Box<dyn Any + Send + Sync> = Box::new(info);	1✔
287		1✔
288	let info3 = info2	1✔
289	.downcast_ref::<Box<dyn MetaData<i32, VectorResultDescriptor, VectorQueryRectangle>>>()	1✔
290	.unwrap();	1✔
291
292	assert_eq!(	1✔
293	info3.result_descriptor().await.unwrap(),	1✔
294	VectorResultDescriptor {	1✔
295	data_type: VectorDataType::Data,	1✔
296	spatial_reference: SpatialReferenceOption::Unreferenced,	1✔
297	columns: Default::default(),	1✔
298	time: None,	1✔
299	bbox: None,	1✔
300	}	1✔
301	);
302	}
303	}

geo-engine / geoengine / 3929938005

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous