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

90.16

/operators/src/source/ogr_source/dataset_iterator.rs

// generated code of `_OgrDatasetIterator` needs this lint for the `peeked` field
#![allow(clippy::option_option)]

use super::{AttributeFilter, CsvHeader, FeaturesProvider, FormatSpecifics, OgrSourceDataset};
use crate::error::{self};
use crate::util::gdal::gdal_open_dataset_ex;
use crate::util::Result;
use gdal::vector::sql::Dialect;
use gdal::vector::{Feature, LayerAccess};
use gdal::{Dataset, DatasetOptions, GdalOpenFlags};
use geoengine_datatypes::primitives::VectorQueryRectangle;
use log::debug;
use ouroboros::self_referencing;
use std::cell::Cell;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::iter::FusedIterator;

/// An iterator over features from a OGR dataset.
/// This iterator contains the dataset and one of its layers.
pub struct OgrDatasetIterator {
    dataset_iterator: _OgrDatasetIterator,
    // must be cell since we borrow self in the iterator for emitting the output value
    // and thus cannot mutably borrow this value
    has_ended: Cell<bool>,
    use_ogr_spatial_filter: bool,
}

// We can implement `Send` for the combination of OGR dataset and layer
// as long we have a one-to-one relation. The layer mutates the dataset.
// So it is not `Send` if there is more than one layer.
unsafe impl Send for OgrDatasetIterator {}

/// Store a dataset and one of its layers.
/// Allows to iterate over features via accessing the layer only.
/// We must ensure to not access it from the outside.
#[self_referencing]
struct _OgrDatasetIterator {
    dataset: gdal::Dataset,
    #[borrows(mut dataset)]
    #[covariant]
    features_provider: FeaturesProvider<'this>,
}

impl OgrDatasetIterator {
    #[allow(clippy::needless_pass_by_value)]
    pub fn new(
        dataset_information: &OgrSourceDataset,
        query_rectangle: &VectorQueryRectangle,
        attribute_filters: Vec<AttributeFilter>,
    ) -> Result<OgrDatasetIterator> {
        let adjusted_filters =
            Self::adjust_filters_to_column_renaming(dataset_information, attribute_filters);

        let dataset_iterator = _OgrDatasetIteratorTryBuilder {
            dataset: Self::open_gdal_dataset(dataset_information)?,
            features_provider_builder: |dataset| {
                Self::create_features_provider(
                    dataset,
                    dataset_information,
                    query_rectangle,
                    &adjusted_filters,
                )
            },
        }
        .try_build()?;

        let use_ogr_spatial_filter = dataset_information.force_ogr_spatial_filter
            || dataset_iterator
                .borrow_features_provider()
                .has_gdal_capability(gdal::vector::LayerCaps::OLCFastSpatialFilter);

        Ok(Self {
            dataset_iterator,
            has_ended: Cell::new(false),
            use_ogr_spatial_filter,
        })
    }

    /// Undo the column renaming to let OGR apply the filters
    fn adjust_filters_to_column_renaming(
        dataset_information: &OgrSourceDataset,
        attribute_filters: Vec<AttributeFilter>,
    ) -> Vec<AttributeFilter> {
        match &dataset_information.columns {
            Some(cspec) => {
                match &cspec.rename {
                    Some(mapping) => {
                        // Build reverse mapping
                        let r_mapping = mapping
                            .iter()
                            .map(|(k, v)| (v.to_string(), k.to_string()))
                            .collect::<HashMap<_, _>>();

                        attribute_filters
                            .into_iter()
                            .map(|f| match r_mapping.get(&f.attribute) {
                                Some(name) => AttributeFilter {
                                    attribute: name.to_string(),
                                    ranges: f.ranges,
                                    keep_nulls: f.keep_nulls,
                                },
                                None => f,
                            })
                            .collect::<Vec<_>>()
                    }
                    // No renaming
                    None => attribute_filters,
                }
            }
            // No column spec
            None => attribute_filters,
        }
    }

    fn create_features_provider<'d>(
        dataset: &'d Dataset,
        dataset_information: &OgrSourceDataset,
        query_rectangle: &VectorQueryRectangle,
        attribute_filters: &[AttributeFilter],
    ) -> Result<FeaturesProvider<'d>> {
        // TODO: add OGR time filter if forced

        let mut features_provider = if let Some(sql) = dataset_information.sql_query.as_ref() {
            FeaturesProvider::ResultSet(
                dataset
                    .execute_sql(sql, None, Dialect::DEFAULT)?
                    .ok_or(error::Error::OgrSqlQuery)?,
            )
        } else {
            FeaturesProvider::Layer(dataset.layer_by_name(&dataset_information.layer_name)?)
        };

        let use_ogr_spatial_filter = dataset_information.force_ogr_spatial_filter
            || features_provider.has_gdal_capability(gdal::vector::LayerCaps::OLCFastSpatialFilter);

        if use_ogr_spatial_filter {
            debug!(
                "using spatial filter {:?} for layer {:?}",
                query_rectangle.spatial_bounds, &dataset_information.layer_name
            );
            // NOTE: the OGR-filter may be inaccurately allowing more features that should be returned in a "strict" fashion.
            features_provider.set_spatial_filter(&query_rectangle.spatial_bounds);
        }

        let filter_string = if dataset.driver().short_name() == "CSV" {
            FeaturesProvider::create_attribute_filter_string_cast(attribute_filters)
        } else {
            FeaturesProvider::create_attribute_filter_string(attribute_filters)
        };

        let final_filter = filter_string
            .map(|f| match &dataset_information.attribute_query {
                Some(a) => format!("({a}) AND {f}"),
                None => f,
            })
            .or_else(|| dataset_information.attribute_query.clone());

        if let Some(filter) = final_filter {
            debug!(
                "using attribute filter {:?} for layer {:?}",
                &filter, &dataset_information.layer_name
            );
            features_provider.set_attribute_filter(filter.as_str())?;
        }
        Ok(features_provider)
    }

    fn open_gdal_dataset(dataset_info: &OgrSourceDataset) -> Result<Dataset> {
        if Self::is_csv(dataset_info) {
            Self::open_csv_dataset(dataset_info)
        } else {
            gdal_open_dataset_ex(
                &dataset_info.file_name,
                DatasetOptions {
                    open_flags: GdalOpenFlags::GDAL_OF_VECTOR,
                    ..Default::default()
                },
            )
        }
    }

    fn open_csv_dataset(dataset_info: &OgrSourceDataset) -> Result<Dataset> {
        let columns = dataset_info
            .columns
            .as_ref()
            .ok_or(error::Error::OgrSourceColumnsSpecMissing)?;

        let allowed_drivers = Some(vec!["CSV"]);

        let mut dataset_options = DatasetOptions {
            open_flags: GdalOpenFlags::GDAL_OF_VECTOR,
            allowed_drivers: allowed_drivers.as_deref(),
            ..DatasetOptions::default()
        };

        let headers = if let Some(FormatSpecifics::Csv { header }) = &columns.format_specifics {
            header.as_gdal_param()
        } else {
            CsvHeader::Auto.as_gdal_param()
        };

        // TODO: make column x optional or allow other indication for data collection
        if columns.x.is_empty() {
            let open_opts = &[
                headers.as_str(),
                // "AUTODETECT_TYPE=YES", // This breaks tests
            ];
            dataset_options.open_options = Some(open_opts);
            return gdal_open_dataset_ex(&dataset_info.file_name, dataset_options);
        }

        if let Some(y) = &columns.y {
            let open_opts = &[
                &format!("X_POSSIBLE_NAMES={}", columns.x),
                &format!("Y_POSSIBLE_NAMES={y}"),
                headers.as_str(),
                "AUTODETECT_TYPE=YES",
            ];
            dataset_options.open_options = Some(open_opts);
            return gdal_open_dataset_ex(&dataset_info.file_name, dataset_options);
        }

        let open_opts = &[
            &format!("GEOM_POSSIBLE_NAMES={}", columns.x),
            headers.as_str(),
            "AUTODETECT_TYPE=YES",
        ];
        dataset_options.open_options = Some(open_opts);
        gdal_open_dataset_ex(&dataset_info.file_name, dataset_options)
    }

    fn is_csv(dataset_info: &OgrSourceDataset) -> bool {
        if let Some("csv" | "tsv") = dataset_info.file_name.extension().and_then(OsStr::to_str) {
            return true;
        }

        dataset_info.file_name.as_path().starts_with("CSV:")
    }

    pub fn was_spatial_filtered_by_ogr(&self) -> bool {
        self.use_ogr_spatial_filter
    }
}

#[allow(clippy::copy_iterator)]
impl<'f> Iterator for &'f mut OgrDatasetIterator {
    type Item = Feature<'f>;

    fn next(&mut self) -> Option<Self::Item> {
        // fuse
        if self.has_ended.get() {
            return None;
        }

        let features_provider = self.dataset_iterator.borrow_features_provider();

        // We somehow have to tell the reference to adhere to the lifetime `'f`
        // On the other hand, we could implement this for `&'f _` instead of `&'f mut _` and get rid of the transmute.
        // However, it makes more sense to require a mutable reference here.
        let features_provider = unsafe { std::mem::transmute::<&'_ _, &'f _>(features_provider) };

        let next = feature_iterator_next(features_provider);

        if next.is_none() {
            self.has_ended.set(true);
        }

        next
    }
}

impl<'f> FusedIterator for &'f mut OgrDatasetIterator {}

// TODO: add this to the `gdal` crate
#[inline]
fn feature_iterator_next<'f>(features_provider: &'f FeaturesProvider) -> Option<Feature<'f>> {
    let layer_ref = features_provider.layer_ref();

    let c_feature = unsafe { gdal_sys::OGR_L_GetNextFeature(layer_ref.c_layer()) };
    if c_feature.is_null() {
        None
    } else {
        Some(unsafe { Feature::from_c_feature(layer_ref.defn(), c_feature) })
    }
}

1	// generated code of `_OgrDatasetIterator` needs this lint for the `peeked` field
2	#![allow(clippy::option_option)]
3
4	use super::{AttributeFilter, CsvHeader, FeaturesProvider, FormatSpecifics, OgrSourceDataset};
5	use crate::error::{self};
6	use crate::util::gdal::gdal_open_dataset_ex;
7	use crate::util::Result;
8	use gdal::vector::sql::Dialect;
9	use gdal::vector::{Feature, LayerAccess};
10	use gdal::{Dataset, DatasetOptions, GdalOpenFlags};
11	use geoengine_datatypes::primitives::VectorQueryRectangle;
12	use log::debug;
13	use ouroboros::self_referencing;
14	use std::cell::Cell;
15	use std::collections::HashMap;
16	use std::ffi::OsStr;
17	use std::iter::FusedIterator;
18
19	/// An iterator over features from a OGR dataset.
20	/// This iterator contains the dataset and one of its layers.
21	pub struct OgrDatasetIterator {
22	dataset_iterator: _OgrDatasetIterator,
23	// must be cell since we borrow self in the iterator for emitting the output value
24	// and thus cannot mutably borrow this value
25	has_ended: Cell<bool>,
26	use_ogr_spatial_filter: bool,
27	}
28
29	// We can implement `Send` for the combination of OGR dataset and layer
30	// as long we have a one-to-one relation. The layer mutates the dataset.
31	// So it is not `Send` if there is more than one layer.
32	unsafe impl Send for OgrDatasetIterator {}
33
34	/// Store a dataset and one of its layers.
35	/// Allows to iterate over features via accessing the layer only.
36	/// We must ensure to not access it from the outside.
37	#[self_referencing]	88✔
38	struct _OgrDatasetIterator {
39	dataset: gdal::Dataset,
40	#[borrows(mut dataset)]
41	#[covariant]
42	features_provider: FeaturesProvider<'this>,
43	}
44
45	impl OgrDatasetIterator {
46	#[allow(clippy::needless_pass_by_value)]
47	pub fn new(	45✔
48	dataset_information: &OgrSourceDataset,	45✔
49	query_rectangle: &VectorQueryRectangle,	45✔
50	attribute_filters: Vec<AttributeFilter>,	45✔
51	) -> Result<OgrDatasetIterator> {	45✔
52	let adjusted_filters =	45✔
53	Self::adjust_filters_to_column_renaming(dataset_information, attribute_filters);	45✔
54
55	let dataset_iterator = _OgrDatasetIteratorTryBuilder {	44✔
56	dataset: Self::open_gdal_dataset(dataset_information)?,	45✔
57	features_provider_builder: \|dataset\| {	44✔
58	Self::create_features_provider(	44✔
59	dataset,	44✔
60	dataset_information,	44✔
61	query_rectangle,	44✔
62	&adjusted_filters,	44✔
63	)	44✔
64	},	44✔
65	}
66	.try_build()?;	44✔
67
68	let use_ogr_spatial_filter = dataset_information.force_ogr_spatial_filter	44✔
69	\|\| dataset_iterator	42✔
70	.borrow_features_provider()	42✔
71	.has_gdal_capability(gdal::vector::LayerCaps::OLCFastSpatialFilter);	42✔
72
73	Ok(Self {	44✔
74	dataset_iterator,	44✔
75	has_ended: Cell::new(false),	44✔
76	use_ogr_spatial_filter,	44✔
77	})	44✔
78	}	45✔
79
80	/// Undo the column renaming to let OGR apply the filters
81	fn adjust_filters_to_column_renaming(	45✔
82	dataset_information: &OgrSourceDataset,	45✔
83	attribute_filters: Vec<AttributeFilter>,	45✔
84	) -> Vec<AttributeFilter> {	45✔
85	match &dataset_information.columns {	45✔
86	Some(cspec) => {	33✔
87	match &cspec.rename {	33✔
88	Some(mapping) => {	7✔
89	// Build reverse mapping	7✔
90	let r_mapping = mapping	7✔
91	.iter()	7✔
92	.map(\|(k, v)\| (v.to_string(), k.to_string()))	191✔
93	.collect::<HashMap<_, _>>();	7✔
94		7✔
95	attribute_filters	7✔
96	.into_iter()	7✔
97	.map(\|f\| match r_mapping.get(&f.attribute) {	7✔
98	Some(name) => AttributeFilter {	1✔
99	attribute: name.to_string(),	1✔
100	ranges: f.ranges,	1✔
101	keep_nulls: f.keep_nulls,	1✔
102	},	1✔
103	None => f,	×
104	})	7✔
105	.collect::<Vec<_>>()	7✔
106	}
107	// No renaming
108	None => attribute_filters,	26✔
109	}
110	}
111	// No column spec
112	None => attribute_filters,	12✔
113	}
114	}	45✔
115
116	fn create_features_provider<'d>(	44✔
117	dataset: &'d Dataset,	44✔
118	dataset_information: &OgrSourceDataset,	44✔
119	query_rectangle: &VectorQueryRectangle,	44✔
120	attribute_filters: &[AttributeFilter],	44✔
121	) -> Result<FeaturesProvider<'d>> {	44✔
122	// TODO: add OGR time filter if forced
123
124	let mut features_provider = if let Some(sql) = dataset_information.sql_query.as_ref() {	44✔
125	FeaturesProvider::ResultSet(
126	dataset	×
127	.execute_sql(sql, None, Dialect::DEFAULT)?	×
128	.ok_or(error::Error::OgrSqlQuery)?,	×
129	)
130	} else {
131	FeaturesProvider::Layer(dataset.layer_by_name(&dataset_information.layer_name)?)	44✔
132	};
133
134	let use_ogr_spatial_filter = dataset_information.force_ogr_spatial_filter	44✔
135	\|\| features_provider.has_gdal_capability(gdal::vector::LayerCaps::OLCFastSpatialFilter);	42✔
136
137	if use_ogr_spatial_filter {	44✔
138	debug!(	4✔
139	"using spatial filter {:?} for layer {:?}",	×
140	query_rectangle.spatial_bounds, &dataset_information.layer_name	×
141	);
142	// NOTE: the OGR-filter may be inaccurately allowing more features that should be returned in a "strict" fashion.
143	features_provider.set_spatial_filter(&query_rectangle.spatial_bounds);	4✔
144	}	40✔
145
146	let filter_string = if dataset.driver().short_name() == "CSV" {	44✔
147	FeaturesProvider::create_attribute_filter_string_cast(attribute_filters)	21✔
148	} else {
149	FeaturesProvider::create_attribute_filter_string(attribute_filters)	23✔
150	};
151
152	let final_filter = filter_string	44✔
153	.map(\|f\| match &dataset_information.attribute_query {	44✔
154	Some(a) => format!("({a}) AND {f}"),	2✔
155	None => f,	10✔
156	})	44✔
157	.or_else(\|\| dataset_information.attribute_query.clone());	44✔
158
159	if let Some(filter) = final_filter {	44✔
160	debug!(	13✔
161	"using attribute filter {:?} for layer {:?}",	×
162	&filter, &dataset_information.layer_name	×
163	);
164	features_provider.set_attribute_filter(filter.as_str())?;	13✔
165	}	31✔
166	Ok(features_provider)	44✔
167	}	44✔
168
169	fn open_gdal_dataset(dataset_info: &OgrSourceDataset) -> Result<Dataset> {	45✔
170	if Self::is_csv(dataset_info) {	45✔
171	Self::open_csv_dataset(dataset_info)	21✔
172	} else {
173	gdal_open_dataset_ex(	24✔
174	&dataset_info.file_name,	24✔
175	DatasetOptions {	24✔
176	open_flags: GdalOpenFlags::GDAL_OF_VECTOR,	24✔
177	..Default::default()	24✔
178	},	24✔
179	)	24✔
180	}
181	}	45✔
182
183	fn open_csv_dataset(dataset_info: &OgrSourceDataset) -> Result<Dataset> {	21✔
184	let columns = dataset_info	21✔
185	.columns	21✔
186	.as_ref()	21✔
187	.ok_or(error::Error::OgrSourceColumnsSpecMissing)?;	21✔
188
189	let allowed_drivers = Some(vec!["CSV"]);	21✔
190		21✔
191	let mut dataset_options = DatasetOptions {	21✔
192	open_flags: GdalOpenFlags::GDAL_OF_VECTOR,	21✔
193	allowed_drivers: allowed_drivers.as_deref(),	21✔
194	..DatasetOptions::default()	21✔
195	};	21✔
196
197	let headers = if let Some(FormatSpecifics::Csv { header }) = &columns.format_specifics {	21✔
198	header.as_gdal_param()	21✔
199	} else {
200	CsvHeader::Auto.as_gdal_param()	×
201	};
202
203	// TODO: make column x optional or allow other indication for data collection
204	if columns.x.is_empty() {	21✔
205	let open_opts = &[	13✔
206	headers.as_str(),	13✔
207	// "AUTODETECT_TYPE=YES", // This breaks tests	13✔
208	];	13✔
209	dataset_options.open_options = Some(open_opts);	13✔
210	return gdal_open_dataset_ex(&dataset_info.file_name, dataset_options);	13✔
211	}	8✔
212
213	if let Some(y) = &columns.y {	8✔
214	let open_opts = &[	8✔
215	&format!("X_POSSIBLE_NAMES={}", columns.x),	8✔
216	&format!("Y_POSSIBLE_NAMES={y}"),	8✔
217	headers.as_str(),	8✔
218	"AUTODETECT_TYPE=YES",	8✔
219	];	8✔
220	dataset_options.open_options = Some(open_opts);	8✔
221	return gdal_open_dataset_ex(&dataset_info.file_name, dataset_options);	8✔
222	}	×
223		×
224	let open_opts = &[	×
225	&format!("GEOM_POSSIBLE_NAMES={}", columns.x),	×
226	headers.as_str(),	×
227	"AUTODETECT_TYPE=YES",	×
228	];	×
229	dataset_options.open_options = Some(open_opts);	×
230	gdal_open_dataset_ex(&dataset_info.file_name, dataset_options)	×
231	}	21✔
232
233	fn is_csv(dataset_info: &OgrSourceDataset) -> bool {
234	if let Some("csv" \| "tsv") = dataset_info.file_name.extension().and_then(OsStr::to_str) {	45✔
235	return true;	17✔
236	}	28✔
237		28✔
238	dataset_info.file_name.as_path().starts_with("CSV:")	28✔
239	}	45✔
240
241	pub fn was_spatial_filtered_by_ogr(&self) -> bool {	232✔
242	self.use_ogr_spatial_filter	232✔
243	}	232✔
244	}
245
246	#[allow(clippy::copy_iterator)]
247	impl<'f> Iterator for &'f mut OgrDatasetIterator {
248	type Item = Feature<'f>;
249
250	fn next(&mut self) -> Option<Self::Item> {	9,716✔
251	// fuse	9,716✔
252	if self.has_ended.get() {	9,716✔
253	return None;	38✔
254	}	9,678✔
255		9,678✔
256	let features_provider = self.dataset_iterator.borrow_features_provider();	9,678✔
257		9,678✔
258	// We somehow have to tell the reference to adhere to the lifetime `'f`	9,678✔
259	// On the other hand, we could implement this for `&'f _` instead of `&'f mut _` and get rid of the transmute.	9,678✔
260	// However, it makes more sense to require a mutable reference here.	9,678✔
261	let features_provider = unsafe { std::mem::transmute::<&'_ _, &'f _>(features_provider) };	9,678✔
262		9,678✔
263	let next = feature_iterator_next(features_provider);	9,678✔
264		9,678✔
265	if next.is_none() {	9,678✔
266	self.has_ended.set(true);	42✔
267	}	9,636✔
268
269	next	9,678✔
270	}	9,716✔
271	}
272
273	impl<'f> FusedIterator for &'f mut OgrDatasetIterator {}
274
275	// TODO: add this to the `gdal` crate
276	#[inline]
277	fn feature_iterator_next<'f>(features_provider: &'f FeaturesProvider) -> Option<Feature<'f>> {	9,678✔
278	let layer_ref = features_provider.layer_ref();	9,678✔
279		9,678✔
280	let c_feature = unsafe { gdal_sys::OGR_L_GetNextFeature(layer_ref.c_layer()) };	9,678✔
281	if c_feature.is_null() {	9,678✔
282	None	42✔
283	} else {
284	Some(unsafe { Feature::from_c_feature(layer_ref.defn(), c_feature) })	9,636✔
285	}
286	}	9,678✔

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