23550417448

Committed 25 Mar 2026 03:52PM UTC coverage: 87.389%. First build

Build # 23550417448

Build Type

Pull #1114

github

Committed by

web-flow

Commit Message

Merge 17fe5e1a3 into dccacebf2

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

Coverage Stats

598 of 1795 new or added lines in 12 files covered. (33.31%)

113704 of 130113 relevant lines covered (87.39%)

497604.14 hits per line

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

90.71

/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 {
                                // fail in tests
                                debug_assert!(false, "BandFilter: received tile for band {} which is not in the selected bands", tile.band);

                                // in production: log warning
                                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 {
251	// fail in tests
NEW 252	debug_assert!(false, "BandFilter: received tile for band {} which is not in the selected bands", tile.band);	×
253
254	// in production: log warning
NEW 255	warn!("BandFilter: Received tile for band {} which is not in the selected bands, skipping. This is a bug.", tile.band);	×
256
NEW 257	None	×
258	}
259	}
NEW 260	Err(e) => Some(Err(e)),	×
261	}
262	}	4✔
263	})	4✔
264	.boxed())
265	}	2✔
266
267	fn result_descriptor(&self) -> &RasterResultDescriptor {	4✔
268	&self.result_descriptor	4✔
269	}	4✔
270	}
271
272	#[async_trait]
273	impl<S> RasterQueryProcessor for BandFilterProcessor<S>
274	where
275	S: RasterQueryProcessor,
276	{
277	type RasterType = S::RasterType;
278
279	async fn _time_query<'a>(
280	&'a self,
281	query: geoengine_datatypes::primitives::TimeInterval,
282	ctx: &'a dyn QueryContext,
NEW 283	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {	×
284	self.source.time_query(query, ctx).await
NEW 285	}	×
286	}
287
288	#[cfg(test)]
289	mod tests {
290	use super::*;
291
292	use crate::{
293	engine::{
294	MockExecutionContext, RasterBandDescriptor, RasterBandDescriptors,
295	SpatialGridDescriptor, TimeDescriptor,
296	},
297	mock::{MockRasterSource, MockRasterSourceParams},
298	};
299	use futures::StreamExt;
300	use geoengine_datatypes::{
301	primitives::{CacheHint, Coordinate2D, TimeInterval},
302	raster::{
303	BoundedGrid, GeoTransform, Grid2D, GridBoundingBox2D, GridShape2D, MaskedGrid2D,
304	RasterDataType, TileInformation, TilesEqualIgnoringCacheHint,
305	},
306	spatial_reference::SpatialReference,
307	util::test::TestDefault,
308	};
309
310	#[tokio::test]
311	#[allow(clippy::too_many_lines)]
312	async fn it_filters_bands() {	1✔
313	let tile_size_in_pixels = GridShape2D::new_2d(2, 2);	1✔
314	let time_interval = TimeInterval::default();	1✔
315
316	let result_descriptor = RasterResultDescriptor {	1✔
317	data_type: RasterDataType::U8,	1✔
318	spatial_reference: SpatialReference::epsg_4326().into(),	1✔
319	time: TimeDescriptor::new_irregular(Some(time_interval)),	1✔
320	spatial_grid: SpatialGridDescriptor::source_from_parts(	1✔
321	GeoTransform::new(Coordinate2D::new(0., 0.), 1., -1.),	1✔
322	tile_size_in_pixels.bounding_box(),	1✔
323	),	1✔
324	bands: RasterBandDescriptors::new(vec![	1✔
325	RasterBandDescriptor::new_unitless("red".into()),	1✔
326	RasterBandDescriptor::new_unitless("green".into()),	1✔
327	RasterBandDescriptor::new_unitless("blue".into()),	1✔
328	])	1✔
329	.unwrap(),	1✔
330	};	1✔
331
332	let tile_info = TileInformation {	1✔
333	global_geo_transform: TestDefault::test_default(),	1✔
334	global_tile_position: [0, 0].into(),	1✔
335	tile_size_in_pixels,	1✔
336	};	1✔
337
338	let band_0: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![1_u8, 2, 3, 4])	1✔
339	.unwrap()	1✔
340	.into();	1✔
341	let band_1: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![5_u8, 6, 7, 8])	1✔
342	.unwrap()	1✔
343	.into();	1✔
344	let band_2: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![9_u8, 10, 11, 12])	1✔
345	.unwrap()	1✔
346	.into();	1✔
347
348	let data = vec![	1✔
349	RasterTile2D::new_with_tile_info(	1✔
350	time_interval,	1✔
351	tile_info,	1✔
352	0,
353	band_0.into(),	1✔
354	CacheHint::default(),	1✔
355	),
356	RasterTile2D::new_with_tile_info(	1✔
357	time_interval,	1✔
358	tile_info,	1✔
359	1,
360	band_1.into(),	1✔
361	CacheHint::default(),	1✔
362	),
363	RasterTile2D::new_with_tile_info(	1✔
364	time_interval,	1✔
365	tile_info,	1✔
366	2,
367	band_2.into(),	1✔
368	CacheHint::default(),	1✔
369	),
370	];
371
372	let mock_source = MockRasterSource {	1✔
373	params: MockRasterSourceParams {	1✔
374	data: data.clone(),	1✔
375	result_descriptor: result_descriptor.clone(),	1✔
376	},	1✔
377	}	1✔
378	.boxed();	1✔
379
380	let band_filter: Box<dyn RasterOperator> = BandFilter {	1✔
381	params: BandFilterParams {	1✔
382	bands: BandsByNameOrIndex::Name(vec!["red".to_string(), "blue".to_string()]),	1✔
383	},	1✔
384	sources: SingleRasterSource {	1✔
385	raster: mock_source,	1✔
386	},	1✔
387	}	1✔
388	.boxed();	1✔
389
390	let execution_context = MockExecutionContext::new_with_tiling_spec(	1✔
391	geoengine_datatypes::raster::TilingSpecification::new(tile_size_in_pixels),	1✔
392	);
393	let query_context = execution_context.mock_query_context_test_default();	1✔
394
395	let query_rect = RasterQueryRectangle::new(	1✔
396	GridBoundingBox2D::new([0, 0], [1, 1]).unwrap(),	1✔
397	time_interval,	1✔
398	[0, 1].try_into().unwrap(),	1✔
399	);
400
401	let initialized = band_filter	1✔
402	.initialize(WorkflowOperatorPath::initialize_root(), &execution_context)	1✔
403	.await	1✔
404	.unwrap();	1✔
405
406	let qp = initialized.query_processor().unwrap().get_u8().unwrap();	1✔
407
408	let result = qp	1✔
409	.raster_query(query_rect, &query_context)	1✔
410	.await	1✔
411	.unwrap()	1✔
412	.collect::<Vec<_>>()	1✔
413	.await;	1✔
414	let result = result.into_iter().collect::<Result<Vec<_>>>().unwrap();	1✔
415
416	// Expected: band 0 (red) and band 1 (blue from original band 2)
417	let mut expected_band_0 = data[0].clone();	1✔
418	expected_band_0.band = 0;	1✔
419	let mut expected_band_1 = data[2].clone();	1✔
420	expected_band_1.band = 1;	1✔
421	let expected = vec![expected_band_0, expected_band_1];	1✔
422
423	assert!(expected.tiles_equal_ignoring_cache_hint(&result));	1✔
424
425	assert_eq!(	1✔
426	initialized.result_descriptor().bands,	1✔
427	vec![	1✔
428	RasterBandDescriptor::new_unitless("red".into()),	1✔
429	RasterBandDescriptor::new_unitless("blue".into()),	1✔
430	]	1✔
431	.try_into()	1✔
432	.unwrap()	1✔
433	);	1✔
434	}	1✔
435
436	#[tokio::test]
437	#[allow(clippy::too_many_lines)]
438	async fn it_filters_bands_by_index() {	1✔
439	let tile_size_in_pixels = GridShape2D::new_2d(2, 2);	1✔
440	let time_interval = TimeInterval::default();	1✔
441
442	let result_descriptor = RasterResultDescriptor {	1✔
443	data_type: RasterDataType::U8,	1✔
444	spatial_reference: SpatialReference::epsg_4326().into(),	1✔
445	time: TimeDescriptor::new_irregular(Some(time_interval)),	1✔
446	spatial_grid: SpatialGridDescriptor::source_from_parts(	1✔
447	GeoTransform::new(Coordinate2D::new(0., 0.), 1., -1.),	1✔
448	tile_size_in_pixels.bounding_box(),	1✔
449	),	1✔
450	bands: RasterBandDescriptors::new(vec![	1✔
451	RasterBandDescriptor::new_unitless("band0".into()),	1✔
452	RasterBandDescriptor::new_unitless("band1".into()),	1✔
453	RasterBandDescriptor::new_unitless("band2".into()),	1✔
454	])	1✔
455	.unwrap(),	1✔
456	};	1✔
457
458	let tile_info = TileInformation {	1✔
459	global_geo_transform: TestDefault::test_default(),	1✔
460	global_tile_position: [0, 0].into(),	1✔
461	tile_size_in_pixels,	1✔
462	};	1✔
463
464	let band_0: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![1_u8, 2, 3, 4])	1✔
465	.unwrap()	1✔
466	.into();	1✔
467	let band_1: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![5_u8, 6, 7, 8])	1✔
468	.unwrap()	1✔
469	.into();	1✔
470	let band_2: MaskedGrid2D<u8> = Grid2D::new(tile_size_in_pixels, vec![9_u8, 10, 11, 12])	1✔
471	.unwrap()	1✔
472	.into();	1✔
473
474	let data = vec![	1✔
475	RasterTile2D::new_with_tile_info(	1✔
476	time_interval,	1✔
477	tile_info,	1✔
478	0,
479	band_0.into(),	1✔
480	CacheHint::default(),	1✔
481	),
482	RasterTile2D::new_with_tile_info(	1✔
483	time_interval,	1✔
484	tile_info,	1✔
485	1,
486	band_1.into(),	1✔
487	CacheHint::default(),	1✔
488	),
489	RasterTile2D::new_with_tile_info(	1✔
490	time_interval,	1✔
491	tile_info,	1✔
492	2,
493	band_2.into(),	1✔
494	CacheHint::default(),	1✔
495	),
496	];
497
498	let mock_source = MockRasterSource {	1✔
499	params: MockRasterSourceParams {	1✔
500	data: data.clone(),	1✔
501	result_descriptor: result_descriptor.clone(),	1✔
502	},	1✔
503	}	1✔
504	.boxed();	1✔
505
506	let band_filter: Box<dyn RasterOperator> = BandFilter {	1✔
507	params: BandFilterParams {	1✔
508	bands: BandsByNameOrIndex::Index(vec![1, 2]),	1✔
509	},	1✔
510	sources: SingleRasterSource {	1✔
511	raster: mock_source,	1✔
512	},	1✔
513	}	1✔
514	.boxed();	1✔
515
516	let execution_context = MockExecutionContext::new_with_tiling_spec(	1✔
517	geoengine_datatypes::raster::TilingSpecification::new(tile_size_in_pixels),	1✔
518	);
519	let query_context = execution_context.mock_query_context_test_default();	1✔
520
521	let query_rect = RasterQueryRectangle::new(	1✔
522	GridBoundingBox2D::new([0, 0], [1, 1]).unwrap(),	1✔
523	time_interval,	1✔
524	[0, 1].try_into().unwrap(),	1✔
525	);
526
527	let initialized = band_filter	1✔
528	.initialize(WorkflowOperatorPath::initialize_root(), &execution_context)	1✔
529	.await	1✔
530	.unwrap();	1✔
531
532	let qp = initialized.query_processor().unwrap().get_u8().unwrap();	1✔
533
534	let result = qp	1✔
535	.raster_query(query_rect, &query_context)	1✔
536	.await	1✔
537	.unwrap()	1✔
538	.collect::<Vec<_>>()	1✔
539	.await;	1✔
540	let result = result.into_iter().collect::<Result<Vec<_>>>().unwrap();	1✔
541
542	// Expected: band 0 (from original band 1) and band 1 (from original band 2)
543	let mut expected_band_0 = data[1].clone();	1✔
544	expected_band_0.band = 0;	1✔
545	let mut expected_band_1 = data[2].clone();	1✔
546	expected_band_1.band = 1;	1✔
547	let expected = vec![expected_band_0, expected_band_1];	1✔
548
549	assert!(expected.tiles_equal_ignoring_cache_hint(&result));	1✔
550
551	assert_eq!(	1✔
552	initialized.result_descriptor().bands,	1✔
553	vec![	1✔
554	RasterBandDescriptor::new_unitless("band1".into()),	1✔
555	RasterBandDescriptor::new_unitless("band2".into()),	1✔
556	]	1✔
557	.try_into()	1✔
558	.unwrap()	1✔
559	);	1✔
560	}	1✔
561
562	#[test]
563	fn test_serde_by_name() {	1✔
564	let params = BandFilterParams {	1✔
565	bands: BandsByNameOrIndex::Name(vec!["red".to_string(), "green".to_string()]),	1✔
566	};	1✔
567
568	let json = serde_json::to_string(&params).unwrap();	1✔
569	assert_eq!(json, r#"{"bands":["red","green"]}"#);	1✔
570
571	let deserialized: BandFilterParams = serde_json::from_str(&json).unwrap();	1✔
572	assert_eq!(params, deserialized);	1✔
573	}	1✔
574
575	#[test]
576	fn test_serde_by_index() {	1✔
577	let params = BandFilterParams {	1✔
578	bands: BandsByNameOrIndex::Index(vec![0, 2, 4]),	1✔
579	};	1✔
580
581	let json = serde_json::to_string(&params).unwrap();	1✔
582	assert_eq!(json, r#"{"bands":[0,2,4]}"#);	1✔
583
584	let deserialized: BandFilterParams = serde_json::from_str(&json).unwrap();	1✔
585	assert_eq!(params, deserialized);	1✔
586	}	1✔
587	}

geo-engine / geoengine / 23550417448

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