23302516318

Committed 19 Mar 2026 03:25PM UTC coverage: 87.49%. First build

Build # 23302516318

Build Type

Pull #1114

github

Committed by

web-flow

Commit Message

Merge 22372bf8e into dccacebf2

Pull Request Pull Request #1114: feat: STAC dataset import

Coverage Stats

302 of 1467 new or added lines in 8 files covered. (20.59%)

113549 of 129785 relevant lines covered (87.49%)

498861.59 hits per line

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

91.04

/operators/src/processing/band_filter.rs

use crate::engine::{
    CanonicOperatorName, ExecutionContext, InitializedRasterOperator, InitializedSources, Operator,
    OperatorName, QueryContext, QueryProcessor, RasterOperator, RasterQueryProcessor,
    RasterResultDescriptor, SingleRasterSource, TypedRasterQueryProcessor, WorkflowOperatorPath,
};
use crate::optimization::OptimizationError;
use crate::util::Result;
use async_trait::async_trait;
use futures::StreamExt;
use futures::stream::BoxStream;
use geoengine_datatypes::primitives::{BandSelection, RasterQueryRectangle, SpatialResolution};
use geoengine_datatypes::raster::{GridBoundingBox2D, RasterTile2D};
use serde::{Deserialize, Serialize};
use snafu::{Snafu, ensure};
use tracing::{debug, warn};

/// Filter bands of a raster
///
/// Removes all but the specified bands from a raster operator output.
/// The order of the remaining bands is preserved.
pub type BandFilter = Operator<BandFilterParams, SingleRasterSource>;

impl OperatorName for BandFilter {
    const TYPE_NAME: &'static str = "BandFilter";
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BandFilterParams {
    bands: BandsByNameOrIndex,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum BandsByNameOrIndex {
    Name(Vec<String>),
    Index(Vec<usize>),
}

impl BandsByNameOrIndex {
    fn is_empty(&self) -> bool {
        match &self {
            BandsByNameOrIndex::Name(names) => names.is_empty(),
            BandsByNameOrIndex::Index(indices) => indices.is_empty(),
        }
    }
}

#[derive(Debug, Snafu)]
#[snafu(visibility(pub(crate)), context(suffix(false)), module(error))]
pub enum BandFilterError {
    #[snafu(display("No bands selected, must select at least one band."))]
    NoBandsSelected,

    #[snafu(display("Band with name '{}' not found in input raster.", name))]
    BandNotFound { name: String },

    #[snafu(display("Bands cannot be mapped to input raster bands."))]
    BandsCannotBeMappedToInput,
}

#[typetag::serde]
#[async_trait]
impl RasterOperator for BandFilter {
    async fn _initialize(
        self: Box<Self>,
        path: WorkflowOperatorPath,
        context: &dyn ExecutionContext,
    ) -> Result<Box<dyn InitializedRasterOperator>> {
        ensure!(!self.params.bands.is_empty(), error::NoBandsSelected);

        let name = CanonicOperatorName::from(&self);

        let initialized_sources = self
            .sources
            .initialize_sources(path.clone(), context)
            .await?;

        debug!("Initializing `BandFilter` with {:?}.", &self.params);

        let input_descriptor = initialized_sources.raster.result_descriptor().clone();

        let input_bands = input_descriptor.bands;

        let mut output_bands = vec![];
        let mut output_band_to_input_band_idx = vec![];

        match &self.params.bands {
            BandsByNameOrIndex::Name(names) => {
                let mut unresolved_names: std::collections::HashSet<_> = names.iter().collect();
                for (band_idx, band) in input_bands.iter().enumerate() {
                    if unresolved_names.remove(&band.name) {
                        output_bands.push(band.clone());
                        output_band_to_input_band_idx.push(band_idx as u32);
                    }
                }
                if let Some(name) = unresolved_names.into_iter().next() {
                    return Err(BandFilterError::BandNotFound { name: name.clone() }.into());
                }
            }
            BandsByNameOrIndex::Index(indices) => {
                let mut unresolved_indices: std::collections::HashSet<_> =
                    indices.iter().copied().collect();
                for (band_idx, band) in input_bands.iter().enumerate() {
                    if unresolved_indices.remove(&band_idx) {
                        output_bands.push(band.clone());
                        output_band_to_input_band_idx.push(band_idx as u32);
                    }
                }
                if let Some(idx) = unresolved_indices.into_iter().next() {
                    return Err(BandFilterError::BandNotFound {
                        name: format!("index {idx}"),
                    }
                    .into());
                }
            }
        }

        let result_descriptor = RasterResultDescriptor {
            data_type: input_descriptor.data_type,
            spatial_reference: input_descriptor.spatial_reference,
            bands: output_bands.try_into()?,
            time: input_descriptor.time,
            spatial_grid: input_descriptor.spatial_grid,
        };

        let initialized_operator = InitializedBandFilter {
            name,
            path,
            source: initialized_sources.raster,
            result_descriptor,
            bands: self.params.bands.clone(),
            output_band_to_input_band_idx,
        };

        Ok(initialized_operator.boxed())
    }

    span_fn!(BandFilter);
}

pub struct InitializedBandFilter {
    name: CanonicOperatorName,
    path: WorkflowOperatorPath,
    source: Box<dyn InitializedRasterOperator>,
    result_descriptor: RasterResultDescriptor,
    bands: BandsByNameOrIndex,
    output_band_to_input_band_idx: Vec<u32>,
}

impl InitializedRasterOperator for InitializedBandFilter {
    fn result_descriptor(&self) -> &RasterResultDescriptor {
        &self.result_descriptor
    }

    fn query_processor(&self) -> Result<TypedRasterQueryProcessor> {
        let source_processor = self.source.query_processor()?;

        Ok(
            call_on_generic_raster_processor!(source_processor, source => BandFilterProcessor {
                source,
                result_descriptor: self.result_descriptor.clone(),
                output_band_to_input_band_idx: self.output_band_to_input_band_idx.clone(),
            }.boxed().into()),
        )
    }

    fn canonic_name(&self) -> CanonicOperatorName {
        self.name.clone()
    }

    fn name(&self) -> &'static str {
        BandFilter::TYPE_NAME
    }

    fn path(&self) -> WorkflowOperatorPath {
        self.path.clone()
    }

    fn optimize(
        &self,
        target_resolution: SpatialResolution,
    ) -> Result<Box<dyn RasterOperator>, OptimizationError> {
        Ok(BandFilter {
            params: BandFilterParams {
                bands: self.bands.clone(),
            },
            sources: SingleRasterSource {
                raster: self.source.optimize(target_resolution)?,
            },
        }
        .boxed())
    }
}

pub struct BandFilterProcessor<S> {
    source: S,
    result_descriptor: RasterResultDescriptor,
    output_band_to_input_band_idx: Vec<u32>,
}

#[async_trait]
impl<S> QueryProcessor for BandFilterProcessor<S>
where
    S: RasterQueryProcessor,
{
    type Output = RasterTile2D<S::RasterType>;
    type ResultDescription = RasterResultDescriptor;
    type Selection = BandSelection;
    type SpatialBounds = GridBoundingBox2D;

    async fn _query<'a>(
        &'a self,
        query: RasterQueryRectangle,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<RasterTile2D<S::RasterType>>>> {
        let output_attributes = query.attributes();

        let input_bands = output_attributes
            .as_slice()
            .iter()
            .map(|band| {
                self.output_band_to_input_band_idx
                    .get(*band as usize)
                    .copied()
            })
            .collect::<Option<Vec<_>>>()
            .ok_or(BandFilterError::BandsCannotBeMappedToInput)?;

        let query = query.select_attributes(BandSelection::new(input_bands.clone())?);

        // Create a mapping from input band index to output band index
        let mut input_to_output_band: std::collections::HashMap<u32, u32> =
            std::collections::HashMap::new();
        for (output_idx, input_idx) in input_bands.iter().enumerate() {
            input_to_output_band.insert(*input_idx, output_idx as u32);
        }

        let stream = self.source.query(query, ctx).await?;

        // Remap band indices in the returned tiles
        Ok(stream
            .filter_map(move |result| {
                let input_to_output_band = input_to_output_band.clone();
                async move {
                    match result {
                        Ok(mut tile) => {
                            if let Some(&output_band) = input_to_output_band.get(&tile.band) {
                                tile.band = output_band;
                                Some(Ok(tile))
                            } else {
                                warn!("BandFilter: Received tile for band {} which is not in the selected bands, skipping. This is a bug.", tile.band);
                                None
                            }
                        }
                        Err(e) => Some(Err(e)),
                    }
                }
            })
            .boxed())
    }

    fn result_descriptor(&self) -> &RasterResultDescriptor {
        &self.result_descriptor
    }
}

#[async_trait]
impl<S> RasterQueryProcessor for BandFilterProcessor<S>
where
    S: RasterQueryProcessor,
{
    type RasterType = S::RasterType;

    async fn _time_query<'a>(
        &'a self,
        query: geoengine_datatypes::primitives::TimeInterval,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {
        self.source.time_query(query, ctx).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::{
        engine::{
            MockExecutionContext, RasterBandDescriptor, RasterBandDescriptors,
            SpatialGridDescriptor, TimeDescriptor,
        },
        mock::{MockRasterSource, MockRasterSourceParams},
    };
    use futures::StreamExt;
    use geoengine_datatypes::{
        primitives::{CacheHint, Coordinate2D, TimeInterval},
        raster::{
            BoundedGrid, GeoTransform, Grid2D, GridBoundingBox2D, GridShape2D, MaskedGrid2D,
            RasterDataType, TileInformation, TilesEqualIgnoringCacheHint,
        },
        spatial_reference::SpatialReference,
        util::test::TestDefault,
    };

    #[tokio::test]
    #[allow(clippy::too_many_lines)]
    async fn it_filters_bands() {
        let tile_size_in_pixels = GridShape2D::new_2d(2, 2);
        let time_interval = TimeInterval::default();

        let result_descriptor = RasterResultDescriptor {
            data_type: RasterDataType::U8,
            spatial_reference: SpatialReference::epsg_4326().into(),
            time: TimeDescriptor::new_irregular(Some(time_interval)),
            spatial_grid: SpatialGridDescriptor::source_from_parts(
                GeoTransform::new(Coordinate2D::new(0., 0.), 1., -1.),
                tile_size_in_pixels.bounding_box(),
            ),
            bands: RasterBandDescriptors::new(vec![
                RasterBandDescriptor::new_unitless("red".into()),
                RasterBandDescriptor::new_unitless("green".into()),
                RasterBandDescriptor::new_unitless("blue".into()),
            ])
            .unwrap(),
        };

        let tile_info = TileInformation {
            global_geo_transform: TestDefault::test_default(),
            global_tile_position: [0, 0].into(),
            tile_size_in_pixels,
        };

        let band_0: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![1_u8, 2, 3, 4])
            .unwrap()
            .into();
        let band_1: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![5_u8, 6, 7, 8])
            .unwrap()
            .into();
        let band_2: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![9_u8, 10, 11, 12])
            .unwrap()
            .into();

        let data = vec![
            RasterTile2D::new_with_tile_info(
                time_interval,
                tile_info,
                0,
                band_0.into(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                time_interval,
                tile_info,
                1,
                band_1.into(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                time_interval,
                tile_info,
                2,
                band_2.into(),
                CacheHint::default(),
            ),
        ];

        let mock_source = MockRasterSource {
            params: MockRasterSourceParams {
                data: data.clone(),
                result_descriptor: result_descriptor.clone(),
            },
        }
        .boxed();

        let band_filter: Box<dyn RasterOperator> = BandFilter {
            params: BandFilterParams {
                bands: BandsByNameOrIndex::Name(vec!["red".to_string(), "blue".to_string()]),
            },
            sources: SingleRasterSource {
                raster: mock_source,
            },
        }
        .boxed();

        let execution_context = MockExecutionContext::new_with_tiling_spec(
            geoengine_datatypes::raster::TilingSpecification::new(tile_size_in_pixels),
        );
        let query_context = execution_context.mock_query_context_test_default();

        let query_rect = RasterQueryRectangle::new(
            GridBoundingBox2D::new([0, 0], [1, 1]).unwrap(),
            time_interval,
            [0, 1].try_into().unwrap(),
        );

        let initialized = band_filter
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap();

        let qp = initialized.query_processor().unwrap().get_u8().unwrap();

        let result = qp
            .raster_query(query_rect, &query_context)
            .await
            .unwrap()
            .collect::<Vec<_>>()
            .await;
        let result = result.into_iter().collect::<Result<Vec<_>>>().unwrap();

        // Expected: band 0 (red) and band 1 (blue from original band 2)
        let mut expected_band_0 = data[0].clone();
        expected_band_0.band = 0;
        let mut expected_band_1 = data[2].clone();
        expected_band_1.band = 1;
        let expected = vec![expected_band_0, expected_band_1];

        assert!(expected.tiles_equal_ignoring_cache_hint(&result));

        assert_eq!(
            initialized.result_descriptor().bands,
            vec![
                RasterBandDescriptor::new_unitless("red".into()),
                RasterBandDescriptor::new_unitless("blue".into()),
            ]
            .try_into()
            .unwrap()
        );
    }

    #[tokio::test]
    #[allow(clippy::too_many_lines)]
    async fn it_filters_bands_by_index() {
        let tile_size_in_pixels = GridShape2D::new_2d(2, 2);
        let time_interval = TimeInterval::default();

        let result_descriptor = RasterResultDescriptor {
            data_type: RasterDataType::U8,
            spatial_reference: SpatialReference::epsg_4326().into(),
            time: TimeDescriptor::new_irregular(Some(time_interval)),
            spatial_grid: SpatialGridDescriptor::source_from_parts(
                GeoTransform::new(Coordinate2D::new(0., 0.), 1., -1.),
                tile_size_in_pixels.bounding_box(),
            ),
            bands: RasterBandDescriptors::new(vec![
                RasterBandDescriptor::new_unitless("band0".into()),
                RasterBandDescriptor::new_unitless("band1".into()),
                RasterBandDescriptor::new_unitless("band2".into()),
            ])
            .unwrap(),
        };

        let tile_info = TileInformation {
            global_geo_transform: TestDefault::test_default(),
            global_tile_position: [0, 0].into(),
            tile_size_in_pixels,
        };

        let band_0: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![1_u8, 2, 3, 4])
            .unwrap()
            .into();
        let band_1: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![5_u8, 6, 7, 8])
            .unwrap()
            .into();
        let band_2: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![9_u8, 10, 11, 12])
            .unwrap()
            .into();

        let data = vec![
            RasterTile2D::new_with_tile_info(
                time_interval,
                tile_info,
                0,
                band_0.into(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                time_interval,
                tile_info,
                1,
                band_1.into(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                time_interval,
                tile_info,
                2,
                band_2.into(),
                CacheHint::default(),
            ),
        ];

        let mock_source = MockRasterSource {
            params: MockRasterSourceParams {
                data: data.clone(),
                result_descriptor: result_descriptor.clone(),
            },
        }
        .boxed();

        let band_filter: Box<dyn RasterOperator> = BandFilter {
            params: BandFilterParams {
                bands: BandsByNameOrIndex::Index(vec![1, 2]),
            },
            sources: SingleRasterSource {
                raster: mock_source,
            },
        }
        .boxed();

        let execution_context = MockExecutionContext::new_with_tiling_spec(
            geoengine_datatypes::raster::TilingSpecification::new(tile_size_in_pixels),
        );
        let query_context = execution_context.mock_query_context_test_default();

        let query_rect = RasterQueryRectangle::new(
            GridBoundingBox2D::new([0, 0], [1, 1]).unwrap(),
            time_interval,
            [0, 1].try_into().unwrap(),
        );

        let initialized = band_filter
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap();

        let qp = initialized.query_processor().unwrap().get_u8().unwrap();

        let result = qp
            .raster_query(query_rect, &query_context)
            .await
            .unwrap()
            .collect::<Vec<_>>()
            .await;
        let result = result.into_iter().collect::<Result<Vec<_>>>().unwrap();

        // Expected: band 0 (from original band 1) and band 1 (from original band 2)
        let mut expected_band_0 = data[1].clone();
        expected_band_0.band = 0;
        let mut expected_band_1 = data[2].clone();
        expected_band_1.band = 1;
        let expected = vec![expected_band_0, expected_band_1];

        assert!(expected.tiles_equal_ignoring_cache_hint(&result));

        assert_eq!(
            initialized.result_descriptor().bands,
            vec![
                RasterBandDescriptor::new_unitless("band1".into()),
                RasterBandDescriptor::new_unitless("band2".into()),
            ]
            .try_into()
            .unwrap()
        );
    }

    #[test]
    fn test_serde_by_name() {
        let params = BandFilterParams {
            bands: BandsByNameOrIndex::Name(vec!["red".to_string(), "green".to_string()]),
        };

        let json = serde_json::to_string(&params).unwrap();
        assert_eq!(json, r#"{"bands":["red","green"]}"#);

        let deserialized: BandFilterParams = serde_json::from_str(&json).unwrap();
        assert_eq!(params, deserialized);
    }

    #[test]
    fn test_serde_by_index() {
        let params = BandFilterParams {
            bands: BandsByNameOrIndex::Index(vec![0, 2, 4]),
        };

        let json = serde_json::to_string(&params).unwrap();
        assert_eq!(json, r#"{"bands":[0,2,4]}"#);

        let deserialized: BandFilterParams = serde_json::from_str(&json).unwrap();
        assert_eq!(params, deserialized);
    }
}

1	use crate::engine::{
2	CanonicOperatorName, ExecutionContext, InitializedRasterOperator, InitializedSources, Operator,
3	OperatorName, QueryContext, QueryProcessor, RasterOperator, RasterQueryProcessor,
4	RasterResultDescriptor, SingleRasterSource, TypedRasterQueryProcessor, WorkflowOperatorPath,
5	};
6	use crate::optimization::OptimizationError;
7	use crate::util::Result;
8	use async_trait::async_trait;
9	use futures::StreamExt;
10	use futures::stream::BoxStream;
11	use geoengine_datatypes::primitives::{BandSelection, RasterQueryRectangle, SpatialResolution};
12	use geoengine_datatypes::raster::{GridBoundingBox2D, RasterTile2D};
13	use serde::{Deserialize, Serialize};
14	use snafu::{Snafu, ensure};
15	use tracing::{debug, warn};
16
17	/// Filter bands of a raster
18	///
19	/// Removes all but the specified bands from a raster operator output.
20	/// The order of the remaining bands is preserved.
21	pub type BandFilter = Operator<BandFilterParams, SingleRasterSource>;
22
23	impl OperatorName for BandFilter {
24	const TYPE_NAME: &'static str = "BandFilter";
25	}
26
27	#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28	pub struct BandFilterParams {
29	bands: BandsByNameOrIndex,
30	}
31
32	#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
33	#[serde(untagged)]
34	pub enum BandsByNameOrIndex {
35	Name(Vec<String>),
36	Index(Vec<usize>),
37	}
38
39	impl BandsByNameOrIndex {
40	fn is_empty(&self) -> bool {	2✔
41	match &self {	2✔
42	BandsByNameOrIndex::Name(names) => names.is_empty(),	1✔
43	BandsByNameOrIndex::Index(indices) => indices.is_empty(),	1✔
44	}
45	}	2✔
46	}
47
48	#[derive(Debug, Snafu)]
49	#[snafu(visibility(pub(crate)), context(suffix(false)), module(error))]
50	pub enum BandFilterError {
51	#[snafu(display("No bands selected, must select at least one band."))]
52	NoBandsSelected,
53
54	#[snafu(display("Band with name '{}' not found in input raster.", name))]
55	BandNotFound { name: String },
56
57	#[snafu(display("Bands cannot be mapped to input raster bands."))]
58	BandsCannotBeMappedToInput,
59	}
60
NEW 61	#[typetag::serde]	×
62	#[async_trait]
63	impl RasterOperator for BandFilter {
64	async fn _initialize(
65	self: Box<Self>,
66	path: WorkflowOperatorPath,
67	context: &dyn ExecutionContext,
68	) -> Result<Box<dyn InitializedRasterOperator>> {	2✔
69	ensure!(!self.params.bands.is_empty(), error::NoBandsSelected);
70
71	let name = CanonicOperatorName::from(&self);
72
73	let initialized_sources = self
74	.sources
75	.initialize_sources(path.clone(), context)
76	.await?;
77
78	debug!("Initializing `BandFilter` with {:?}.", &self.params);
79
80	let input_descriptor = initialized_sources.raster.result_descriptor().clone();
81
82	let input_bands = input_descriptor.bands;
83
84	let mut output_bands = vec![];
85	let mut output_band_to_input_band_idx = vec![];
86
87	match &self.params.bands {
88	BandsByNameOrIndex::Name(names) => {
89	let mut unresolved_names: std::collections::HashSet<_> = names.iter().collect();
90	for (band_idx, band) in input_bands.iter().enumerate() {
91	if unresolved_names.remove(&band.name) {
92	output_bands.push(band.clone());
93	output_band_to_input_band_idx.push(band_idx as u32);
94	}
95	}
96	if let Some(name) = unresolved_names.into_iter().next() {
97	return Err(BandFilterError::BandNotFound { name: name.clone() }.into());
98	}
99	}
100	BandsByNameOrIndex::Index(indices) => {
101	let mut unresolved_indices: std::collections::HashSet<_> =
102	indices.iter().copied().collect();
103	for (band_idx, band) in input_bands.iter().enumerate() {
104	if unresolved_indices.remove(&band_idx) {
105	output_bands.push(band.clone());
106	output_band_to_input_band_idx.push(band_idx as u32);
107	}
108	}
109	if let Some(idx) = unresolved_indices.into_iter().next() {
110	return Err(BandFilterError::BandNotFound {
111	name: format!("index {idx}"),
112	}
113	.into());
114	}
115	}
116	}
117
118	let result_descriptor = RasterResultDescriptor {
119	data_type: input_descriptor.data_type,
120	spatial_reference: input_descriptor.spatial_reference,
121	bands: output_bands.try_into()?,
122	time: input_descriptor.time,
123	spatial_grid: input_descriptor.spatial_grid,
124	};
125
126	let initialized_operator = InitializedBandFilter {
127	name,
128	path,
129	source: initialized_sources.raster,
130	result_descriptor,
131	bands: self.params.bands.clone(),
132	output_band_to_input_band_idx,
133	};
134
135	Ok(initialized_operator.boxed())
136	}	2✔
137
138	span_fn!(BandFilter);
139	}
140
141	pub struct InitializedBandFilter {
142	name: CanonicOperatorName,
143	path: WorkflowOperatorPath,
144	source: Box<dyn InitializedRasterOperator>,
145	result_descriptor: RasterResultDescriptor,
146	bands: BandsByNameOrIndex,
147	output_band_to_input_band_idx: Vec<u32>,
148	}
149
150	impl InitializedRasterOperator for InitializedBandFilter {
151	fn result_descriptor(&self) -> &RasterResultDescriptor {	2✔
152	&self.result_descriptor	2✔
153	}	2✔
154
155	fn query_processor(&self) -> Result<TypedRasterQueryProcessor> {	2✔
156	let source_processor = self.source.query_processor()?;	2✔
157
158	Ok(
159	call_on_generic_raster_processor!(source_processor, source => BandFilterProcessor {	2✔
160	source,	2✔
161	result_descriptor: self.result_descriptor.clone(),	2✔
162	output_band_to_input_band_idx: self.output_band_to_input_band_idx.clone(),	2✔
163	}.boxed().into()),	2✔
164	)
165	}	2✔
166
NEW 167	fn canonic_name(&self) -> CanonicOperatorName {	×
NEW 168	self.name.clone()	×
NEW 169	}	×
170
NEW 171	fn name(&self) -> &'static str {	×
NEW 172	BandFilter::TYPE_NAME	×
NEW 173	}	×
174
NEW 175	fn path(&self) -> WorkflowOperatorPath {	×
NEW 176	self.path.clone()	×
NEW 177	}	×
178
NEW 179	fn optimize(	×
NEW 180	&self,	×
NEW 181	target_resolution: SpatialResolution,	×
NEW 182	) -> Result<Box<dyn RasterOperator>, OptimizationError> {	×
183	Ok(BandFilter {
NEW 184	params: BandFilterParams {	×
NEW 185	bands: self.bands.clone(),	×
NEW 186	},	×
187	sources: SingleRasterSource {
NEW 188	raster: self.source.optimize(target_resolution)?,	×
189	},
190	}
NEW 191	.boxed())	×
NEW 192	}	×
193	}
194
195	pub struct BandFilterProcessor<S> {
196	source: S,
197	result_descriptor: RasterResultDescriptor,
198	output_band_to_input_band_idx: Vec<u32>,
199	}
200
201	#[async_trait]
202	impl<S> QueryProcessor for BandFilterProcessor<S>
203	where
204	S: RasterQueryProcessor,
205	{
206	type Output = RasterTile2D<S::RasterType>;
207	type ResultDescription = RasterResultDescriptor;
208	type Selection = BandSelection;
209	type SpatialBounds = GridBoundingBox2D;
210
211	async fn _query<'a>(
212	&'a self,
213	query: RasterQueryRectangle,
214	ctx: &'a dyn QueryContext,
215	) -> Result<BoxStream<'a, Result<RasterTile2D<S::RasterType>>>> {	2✔
216	let output_attributes = query.attributes();
217
218	let input_bands = output_attributes
219	.as_slice()
220	.iter()
221	.map(\|band\| {	4✔
222	self.output_band_to_input_band_idx	4✔
223	.get(*band as usize)	4✔
224	.copied()	4✔
225	})	4✔
226	.collect::<Option<Vec<_>>>()
227	.ok_or(BandFilterError::BandsCannotBeMappedToInput)?;
228
229	let query = query.select_attributes(BandSelection::new(input_bands.clone())?);
230
231	// Create a mapping from input band index to output band index
232	let mut input_to_output_band: std::collections::HashMap<u32, u32> =
233	std::collections::HashMap::new();
234	for (output_idx, input_idx) in input_bands.iter().enumerate() {
235	input_to_output_band.insert(*input_idx, output_idx as u32);
236	}
237
238	let stream = self.source.query(query, ctx).await?;
239
240	// Remap band indices in the returned tiles
241	Ok(stream
242	.filter_map(move \|result\| {	4✔
243	let input_to_output_band = input_to_output_band.clone();	4✔
244	async move {	4✔
245	match result {	4✔
246	Ok(mut tile) => {	4✔
247	if let Some(&output_band) = input_to_output_band.get(&tile.band) {	4✔
248	tile.band = output_band;	4✔
249	Some(Ok(tile))	4✔
250	} else {
NEW 251	warn!("BandFilter: Received tile for band {} which is not in the selected bands, skipping. This is a bug.", tile.band);	×
NEW 252	None	×
253	}
254	}
NEW 255	Err(e) => Some(Err(e)),	×
256	}
257	}	4✔
258	})	4✔
259	.boxed())
260	}	2✔
261
262	fn result_descriptor(&self) -> &RasterResultDescriptor {	4✔
263	&self.result_descriptor	4✔
264	}	4✔
265	}
266
267	#[async_trait]
268	impl<S> RasterQueryProcessor for BandFilterProcessor<S>
269	where
270	S: RasterQueryProcessor,
271	{
272	type RasterType = S::RasterType;
273
274	async fn _time_query<'a>(
275	&'a self,
276	query: geoengine_datatypes::primitives::TimeInterval,
277	ctx: &'a dyn QueryContext,
NEW 278	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {	×
279	self.source.time_query(query, ctx).await
NEW 280	}	×
281	}
282
283	#[cfg(test)]
284	mod tests {
285	use super::*;
286
287	use crate::{
288	engine::{
289	MockExecutionContext, RasterBandDescriptor, RasterBandDescriptors,
290	SpatialGridDescriptor, TimeDescriptor,
291	},
292	mock::{MockRasterSource, MockRasterSourceParams},
293	};
294	use futures::StreamExt;
295	use geoengine_datatypes::{
296	primitives::{CacheHint, Coordinate2D, TimeInterval},
297	raster::{
298	BoundedGrid, GeoTransform, Grid2D, GridBoundingBox2D, GridShape2D, MaskedGrid2D,
299	RasterDataType, TileInformation, TilesEqualIgnoringCacheHint,
300	},
301	spatial_reference::SpatialReference,
302	util::test::TestDefault,
303	};
304
305	#[tokio::test]
306	#[allow(clippy::too_many_lines)]
307	async fn it_filters_bands() {	1✔
308	let tile_size_in_pixels = GridShape2D::new_2d(2, 2);	1✔
309	let time_interval = TimeInterval::default();	1✔
310
311	let result_descriptor = RasterResultDescriptor {	1✔
312	data_type: RasterDataType::U8,	1✔
313	spatial_reference: SpatialReference::epsg_4326().into(),	1✔
314	time: TimeDescriptor::new_irregular(Some(time_interval)),	1✔
315	spatial_grid: SpatialGridDescriptor::source_from_parts(	1✔
316	GeoTransform::new(Coordinate2D::new(0., 0.), 1., -1.),	1✔
317	tile_size_in_pixels.bounding_box(),	1✔
318	),	1✔
319	bands: RasterBandDescriptors::new(vec![	1✔
320	RasterBandDescriptor::new_unitless("red".into()),	1✔
321	RasterBandDescriptor::new_unitless("green".into()),	1✔
322	RasterBandDescriptor::new_unitless("blue".into()),	1✔
323	])	1✔
324	.unwrap(),	1✔
325	};	1✔
326
327	let tile_info = TileInformation {	1✔
328	global_geo_transform: TestDefault::test_default(),	1✔
329	global_tile_position: [0, 0].into(),	1✔
330	tile_size_in_pixels,	1✔
331	};	1✔
332
333	let band_0: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![1_u8, 2, 3, 4])	1✔
334	.unwrap()	1✔
335	.into();	1✔
336	let band_1: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![5_u8, 6, 7, 8])	1✔
337	.unwrap()	1✔
338	.into();	1✔
339	let band_2: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![9_u8, 10, 11, 12])	1✔
340	.unwrap()	1✔
341	.into();	1✔
342
343	let data = vec![	1✔
344	RasterTile2D::new_with_tile_info(	1✔
345	time_interval,	1✔
346	tile_info,	1✔
347	0,
348	band_0.into(),	1✔
349	CacheHint::default(),	1✔
350	),
351	RasterTile2D::new_with_tile_info(	1✔
352	time_interval,	1✔
353	tile_info,	1✔
354	1,
355	band_1.into(),	1✔
356	CacheHint::default(),	1✔
357	),
358	RasterTile2D::new_with_tile_info(	1✔
359	time_interval,	1✔
360	tile_info,	1✔
361	2,
362	band_2.into(),	1✔
363	CacheHint::default(),	1✔
364	),
365	];
366
367	let mock_source = MockRasterSource {	1✔
368	params: MockRasterSourceParams {	1✔
369	data: data.clone(),	1✔
370	result_descriptor: result_descriptor.clone(),	1✔
371	},	1✔
372	}	1✔
373	.boxed();	1✔
374
375	let band_filter: Box<dyn RasterOperator> = BandFilter {	1✔
376	params: BandFilterParams {	1✔
377	bands: BandsByNameOrIndex::Name(vec!["red".to_string(), "blue".to_string()]),	1✔
378	},	1✔
379	sources: SingleRasterSource {	1✔
380	raster: mock_source,	1✔
381	},	1✔
382	}	1✔
383	.boxed();	1✔
384
385	let execution_context = MockExecutionContext::new_with_tiling_spec(	1✔
386	geoengine_datatypes::raster::TilingSpecification::new(tile_size_in_pixels),	1✔
387	);
388	let query_context = execution_context.mock_query_context_test_default();	1✔
389
390	let query_rect = RasterQueryRectangle::new(	1✔
391	GridBoundingBox2D::new([0, 0], [1, 1]).unwrap(),	1✔
392	time_interval,	1✔
393	[0, 1].try_into().unwrap(),	1✔
394	);
395
396	let initialized = band_filter	1✔
397	.initialize(WorkflowOperatorPath::initialize_root(), &execution_context)	1✔
398	.await	1✔
399	.unwrap();	1✔
400
401	let qp = initialized.query_processor().unwrap().get_u8().unwrap();	1✔
402
403	let result = qp	1✔
404	.raster_query(query_rect, &query_context)	1✔
405	.await	1✔
406	.unwrap()	1✔
407	.collect::<Vec<_>>()	1✔
408	.await;	1✔
409	let result = result.into_iter().collect::<Result<Vec<_>>>().unwrap();	1✔
410
411	// Expected: band 0 (red) and band 1 (blue from original band 2)
412	let mut expected_band_0 = data[0].clone();	1✔
413	expected_band_0.band = 0;	1✔
414	let mut expected_band_1 = data[2].clone();	1✔
415	expected_band_1.band = 1;	1✔
416	let expected = vec![expected_band_0, expected_band_1];	1✔
417
418	assert!(expected.tiles_equal_ignoring_cache_hint(&result));	1✔
419
420	assert_eq!(	1✔
421	initialized.result_descriptor().bands,	1✔
422	vec![	1✔
423	RasterBandDescriptor::new_unitless("red".into()),	1✔
424	RasterBandDescriptor::new_unitless("blue".into()),	1✔
425	]	1✔
426	.try_into()	1✔
427	.unwrap()	1✔
428	);	1✔
429	}	1✔
430
431	#[tokio::test]
432	#[allow(clippy::too_many_lines)]
433	async fn it_filters_bands_by_index() {	1✔
434	let tile_size_in_pixels = GridShape2D::new_2d(2, 2);	1✔
435	let time_interval = TimeInterval::default();	1✔
436
437	let result_descriptor = RasterResultDescriptor {	1✔
438	data_type: RasterDataType::U8,	1✔
439	spatial_reference: SpatialReference::epsg_4326().into(),	1✔
440	time: TimeDescriptor::new_irregular(Some(time_interval)),	1✔
441	spatial_grid: SpatialGridDescriptor::source_from_parts(	1✔
442	GeoTransform::new(Coordinate2D::new(0., 0.), 1., -1.),	1✔
443	tile_size_in_pixels.bounding_box(),	1✔
444	),	1✔
445	bands: RasterBandDescriptors::new(vec![	1✔
446	RasterBandDescriptor::new_unitless("band0".into()),	1✔
447	RasterBandDescriptor::new_unitless("band1".into()),	1✔
448	RasterBandDescriptor::new_unitless("band2".into()),	1✔
449	])	1✔
450	.unwrap(),	1✔
451	};	1✔
452
453	let tile_info = TileInformation {	1✔
454	global_geo_transform: TestDefault::test_default(),	1✔
455	global_tile_position: [0, 0].into(),	1✔
456	tile_size_in_pixels,	1✔
457	};	1✔
458
459	let band_0: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![1_u8, 2, 3, 4])	1✔
460	.unwrap()	1✔
461	.into();	1✔
462	let band_1: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![5_u8, 6, 7, 8])	1✔
463	.unwrap()	1✔
464	.into();	1✔
465	let band_2: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![9_u8, 10, 11, 12])	1✔
466	.unwrap()	1✔
467	.into();	1✔
468
469	let data = vec![	1✔
470	RasterTile2D::new_with_tile_info(	1✔
471	time_interval,	1✔
472	tile_info,	1✔
473	0,
474	band_0.into(),	1✔
475	CacheHint::default(),	1✔
476	),
477	RasterTile2D::new_with_tile_info(	1✔
478	time_interval,	1✔
479	tile_info,	1✔
480	1,
481	band_1.into(),	1✔
482	CacheHint::default(),	1✔
483	),
484	RasterTile2D::new_with_tile_info(	1✔
485	time_interval,	1✔
486	tile_info,	1✔
487	2,
488	band_2.into(),	1✔
489	CacheHint::default(),	1✔
490	),
491	];
492
493	let mock_source = MockRasterSource {	1✔
494	params: MockRasterSourceParams {	1✔
495	data: data.clone(),	1✔
496	result_descriptor: result_descriptor.clone(),	1✔
497	},	1✔
498	}	1✔
499	.boxed();	1✔
500
501	let band_filter: Box<dyn RasterOperator> = BandFilter {	1✔
502	params: BandFilterParams {	1✔
503	bands: BandsByNameOrIndex::Index(vec![1, 2]),	1✔
504	},	1✔
505	sources: SingleRasterSource {	1✔
506	raster: mock_source,	1✔
507	},	1✔
508	}	1✔
509	.boxed();	1✔
510
511	let execution_context = MockExecutionContext::new_with_tiling_spec(	1✔
512	geoengine_datatypes::raster::TilingSpecification::new(tile_size_in_pixels),	1✔
513	);
514	let query_context = execution_context.mock_query_context_test_default();	1✔
515
516	let query_rect = RasterQueryRectangle::new(	1✔
517	GridBoundingBox2D::new([0, 0], [1, 1]).unwrap(),	1✔
518	time_interval,	1✔
519	[0, 1].try_into().unwrap(),	1✔
520	);
521
522	let initialized = band_filter	1✔
523	.initialize(WorkflowOperatorPath::initialize_root(), &execution_context)	1✔
524	.await	1✔
525	.unwrap();	1✔
526
527	let qp = initialized.query_processor().unwrap().get_u8().unwrap();	1✔
528
529	let result = qp	1✔
530	.raster_query(query_rect, &query_context)	1✔
531	.await	1✔
532	.unwrap()	1✔
533	.collect::<Vec<_>>()	1✔
534	.await;	1✔
535	let result = result.into_iter().collect::<Result<Vec<_>>>().unwrap();	1✔
536
537	// Expected: band 0 (from original band 1) and band 1 (from original band 2)
538	let mut expected_band_0 = data[1].clone();	1✔
539	expected_band_0.band = 0;	1✔
540	let mut expected_band_1 = data[2].clone();	1✔
541	expected_band_1.band = 1;	1✔
542	let expected = vec![expected_band_0, expected_band_1];	1✔
543
544	assert!(expected.tiles_equal_ignoring_cache_hint(&result));	1✔
545
546	assert_eq!(	1✔
547	initialized.result_descriptor().bands,	1✔
548	vec![	1✔
549	RasterBandDescriptor::new_unitless("band1".into()),	1✔
550	RasterBandDescriptor::new_unitless("band2".into()),	1✔
551	]	1✔
552	.try_into()	1✔
553	.unwrap()	1✔
554	);	1✔
555	}	1✔
556
557	#[test]
558	fn test_serde_by_name() {	1✔
559	let params = BandFilterParams {	1✔
560	bands: BandsByNameOrIndex::Name(vec!["red".to_string(), "green".to_string()]),	1✔
561	};	1✔
562
563	let json = serde_json::to_string(&params).unwrap();	1✔
564	assert_eq!(json, r#"{"bands":["red","green"]}"#);	1✔
565
566	let deserialized: BandFilterParams = serde_json::from_str(&json).unwrap();	1✔
567	assert_eq!(params, deserialized);	1✔
568	}	1✔
569
570	#[test]
571	fn test_serde_by_index() {	1✔
572	let params = BandFilterParams {	1✔
573	bands: BandsByNameOrIndex::Index(vec![0, 2, 4]),	1✔
574	};	1✔
575
576	let json = serde_json::to_string(&params).unwrap();	1✔
577	assert_eq!(json, r#"{"bands":[0,2,4]}"#);	1✔
578
579	let deserialized: BandFilterParams = serde_json::from_str(&json).unwrap();	1✔
580	assert_eq!(params, deserialized);	1✔
581	}	1✔
582	}

geo-engine / geoengine / 23302516318

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