Merge branch 'main' into workspace-dependencies

Author: jdroenner

datatypes/src/primitives/db_types.rs

@@ -257,7 +257,7 @@ impl From<&RasterColorizer> for RasterColorizerDbType {
                 blue_band,
                 rgb_params: rgba_params,
             } => Self {
-                r#type: RasterColorizerDbTypeType::SingleBand,
+                r#type: RasterColorizerDbTypeType::MultiBand,
                 band: None,
                 band_colorizer: None,
                 red_band: Some(i64::from(*red_band)),

operators/src/util/raster_stream_to_png.rs

@@ -8,9 +8,7 @@ use geoengine_datatypes::operations::image::{ColorMapper, RgbParams};
 use geoengine_datatypes::raster::{FromIndexFn, GridIndexAccess, GridShapeAccess};
 use geoengine_datatypes::{
     operations::image::{Colorizer, RasterColorizer, RgbaColor, ToPng},
-    primitives::{
-        AxisAlignedRectangle, BandSelection, CacheHint, RasterQueryRectangle, TimeInterval,
-    },
+    primitives::{AxisAlignedRectangle, CacheHint, RasterQueryRectangle, TimeInterval},
     raster::{Blit, ConvertDataType, EmptyGrid2D, GeoTransform, GridOrEmpty, Pixel, RasterTile2D},
 };
 use num_traits::AsPrimitive;
@@ -23,14 +21,24 @@ use tracing::{span, Level};
 #[allow(clippy::too_many_arguments)]
 pub async fn raster_stream_to_png_bytes<T: Pixel, C: QueryContext + 'static>(
     processor: Box<dyn RasterQueryProcessor<RasterType = T>>,
-    mut query_rect: RasterQueryRectangle,
+    query_rect: RasterQueryRectangle,
     mut query_ctx: C,
     width: u32,
     height: u32,
     time: Option<TimeInterval>,
     raster_colorizer: Option<RasterColorizer>,
     conn_closed: BoxFuture<'_, ()>,
 ) -> Result<(Vec<u8>, CacheHint)> {
+    debug_assert!(
+        query_rect.attributes.count() <= 3
+            || query_rect
+                .attributes
+                .as_slice()
+                .windows(2)
+                .all(|w| w[0] < w[1]), // TODO: replace with `is_sorted` once it is stable
+        "bands must be sorted and at most three bands can be queried"
+    );
+
     let span = span!(Level::TRACE, "raster_stream_to_png_bytes");
     let _enter = span.enter();
 
@@ -52,21 +60,24 @@ pub async fn raster_stream_to_png_bytes<T: Pixel, C: QueryContext + 'static>(
         RasterColorizer::SingleBand { band, .. } => vec![band],
     };
 
-    if query_rect
-        .attributes
-        .as_slice()
+    let band_positions = required_bands
         .iter()
-        .any(|band| !required_bands.contains(band))
-    {
+        .filter_map(|band| {
+            query_rect
+                .attributes
+                .as_slice()
+                .iter()
+                .position(|b| b == band)
+        })
+        .collect::<Vec<usize>>();
+
+    if band_positions.len() != required_bands.len() {
         return Err(PngCreationError::ColorizerBandsMustBePresentInQuery {
             bands_present: query_rect.attributes.as_vec(),
             required_bands,
         })?;
     }
 
-    // modify query to only query the required bands
-    query_rect.attributes = BandSelection::new_unchecked(required_bands);
-
     let query_abort_trigger = query_ctx.abort_trigger()?;
 
     let x_query_resolution = query_rect.spatial_bounds.size_x() / f64::from(width);
@@ -114,6 +125,7 @@ pub async fn raster_stream_to_png_bytes<T: Pixel, C: QueryContext + 'static>(
                 width,
                 height,
                 rgba_params,
+                band_positions,
                 conn_closed,
                 query_abort_trigger,
             )
@@ -160,45 +172,45 @@ async fn multi_band_colorizer_to_png_bytes<T: Pixel, C: QueryContext + 'static>(
     width: u32,
     height: u32,
     rgb_params: RgbParams,
+    band_positions: Vec<usize>,
     conn_closed: BoxFuture<'_, ()>,
     query_abort_trigger: QueryAbortTrigger,
 ) -> Result<(Vec<u8>, CacheHint)> {
-    const RGB_CHANNEL_COUNT: usize = 3;
-
-    debug_assert_eq!(query_rect.attributes.count(), RGB_CHANNEL_COUNT as u32);
-
-    let tile_stream = processor.query(query_rect.clone(), &query_ctx).await?;
-
-    let band_numbers: [u32; 3] = [
-        query_rect.attributes.as_slice()[0],
-        query_rect.attributes.as_slice()[1],
-        query_rect.attributes.as_slice()[2],
-    ];
-
+    let rgb_channel_count = query_rect.attributes.count() as usize;
     let no_data_color = rgb_params.no_data_color;
-
     let tile_template: RasterTile2D<u32> = tile_template.convert_data_type();
-    let output_tile = Box::pin(
-        tile_stream
-            // .map_ok(|tile| tile.convert_data_type())
-            .try_chunks(RGB_CHANNEL_COUNT)
-            .fold(Ok(tile_template), |raster2d, chunk| async move {
-                let chunk = chunk.boxed_context(error::QueryDidNotProduceNextChunk)?;
+    let red_band_index = band_positions[0];
+    let green_band_index = band_positions[1];
+    let blue_band_index = band_positions[2];
 
-                let mut tuple: [RasterTile2D<T>; RGB_CHANNEL_COUNT] = chunk
-                    .try_into()
-                    .map_err(|_| PngCreationError::RgbChunkIsNotThreeBands)?;
+    let tile_stream = processor.query(query_rect.clone(), &query_ctx).await?;
 
-                // tiles don't arrive in query order, so we need to sort them to be in query order
-                // TODO: remove this when we have a guarantee that tiles are in query order
-                sort_by_indices(&mut tuple, &band_numbers);
+    let output_tile = Box::pin(tile_stream.try_chunks(rgb_channel_count).fold(
+        Ok(tile_template),
+        |raster2d, chunk| async move {
+            let chunk = chunk.boxed_context(error::QueryDidNotProduceNextChunk)?;
 
-                // TODO: spawn blocking task
-                let rgb_tile = compute_rgb_tile(tuple, &rgb_params);
+            if chunk.len() != rgb_channel_count {
+                return Err(PngCreationError::RgbChunkIsNotEnoughBands)?;
+            }
 
-                blit_tile(raster2d, Ok(rgb_tile))
-            }),
-    );
+            // TODO: spawn blocking task
+            let rgb_tile = crate::util::spawn_blocking(move || {
+                compute_rgb_tile(
+                    [
+                        &chunk[red_band_index],
+                        &chunk[green_band_index],
+                        &chunk[blue_band_index],
+                    ],
+                    &rgb_params,
+                )
+            })
+            .await
+            .boxed_context(error::UnexpectedComputational)?;
+
+            blit_tile(raster2d, Ok(rgb_tile))
+        },
+    ));
 
     let result = abortable_query_execution(output_tile, conn_closed, query_abort_trigger).await?;
     Ok((
@@ -209,24 +221,6 @@ async fn multi_band_colorizer_to_png_bytes<T: Pixel, C: QueryContext + 'static>(
     ))
 }
 
-/// Sorts the tiles by their band numbers and returns them in the order of the band numbers.
-/// Uses bubble sort since the number of bands is small.
-fn sort_by_indices<T>(tiles: &mut [T; 3], band_numbers: &[u32; 3]) {
-    debug_assert_eq!(tiles.len(), 3);
-    debug_assert_eq!(band_numbers.len(), 3);
-
-    let mut band_numbers = [band_numbers[0], band_numbers[1], band_numbers[2]];
-
-    for r in [2, 1] {
-        for i in 0..r {
-            if band_numbers[i] > band_numbers[i + 1] {
-                tiles.swap(i, i + 1);
-                band_numbers.swap(i, i + 1);
-            }
-        }
-    }
-}
-
 fn blit_tile<T>(
     raster2d: Result<RasterTile2D<T>>,
     tile: Result<RasterTile2D<T>>,
@@ -271,7 +265,7 @@ pub fn default_colorizer_gradient<T: Pixel>() -> Colorizer {
 }
 
 pub fn compute_rgb_tile<P: Pixel>(
-    [red, green, blue]: [RasterTile2D<P>; 3],
+    [red, green, blue]: [&RasterTile2D<P>; 3],
     params: &RgbParams,
 ) -> RasterTile2D<u32> {
     fn fit_to_interval_0_255(value: f64, min: f64, max: f64, scale: f64) -> u32 {
@@ -347,8 +341,10 @@ pub enum PngCreationError {
     },
     #[snafu(display("Query did not produce next chunk: {source}"))]
     QueryDidNotProduceNextChunk { source: Box<dyn ErrorSource> },
-    #[snafu(display("RGB chunk is not three bands"))]
-    RgbChunkIsNotThreeBands,
+    #[snafu(display("RGB chunk is not enough bands"))]
+    RgbChunkIsNotEnoughBands,
+    #[snafu(display("Unexpected computational error"))]
+    UnexpectedComputationalError { source: Box<dyn ErrorSource> },
 }
 
 #[cfg(test)]
@@ -429,21 +425,4 @@ mod tests {
             image_bytes.as_slice()
         );
     }
-
-    #[test]
-    fn it_sorts_by_index() {
-        let mut values = [2, 1, 3];
-        let band_numbers = [2, 1, 5];
-
-        sort_by_indices(&mut values, &band_numbers);
-
-        assert_eq!(values, [1, 2, 3]);
-
-        let mut values = [3, 2, 1];
-        let band_numbers = [5, 2, 1];
-
-        sort_by_indices(&mut values, &band_numbers);
-
-        assert_eq!(values, [1, 2, 3]);
-    }
 }

services/src/api/handlers/wms.rs

@@ -348,8 +348,8 @@ async fn wms_map_handler<C: ApplicationContext>(
 
         let attributes = raster_colorizer.as_ref().map_or_else(
             || BandSelection::new_single(0),
-            |o| {
-                RasterColorizer::band_selection(o)
+            |colorizer: &RasterColorizer| {
+                RasterColorizer::band_selection(colorizer)
                     .try_into()
                     .expect("conversion of usize to u32 succeeds for small band numbers")
             },
@@ -482,7 +482,8 @@ mod tests {
     use crate::ge_context;
     use crate::util::tests::{
         check_allowed_http_methods, read_body_string, register_ndvi_workflow_helper,
-        register_ndvi_workflow_helper_with_cache_ttl, send_test_request,
+        register_ndvi_workflow_helper_with_cache_ttl, register_ne2_multiband_workflow,
+        send_test_request,
     };
     use actix_http::header::{self, CONTENT_TYPE};
     use actix_web::dev::ServiceResponse;
@@ -861,6 +862,143 @@ mod tests {
         );
     }
 
+    #[ge_context::test(tiling_spec = "get_map_test_helper_tiling_spec")]
+    async fn it_supports_multiband_colorizer(app_ctx: PostgresContext<NoTls>) {
+        let ctx = app_ctx.default_session_context().await.unwrap();
+
+        let session_id = ctx.session().id();
+
+        let (_, id) = register_ne2_multiband_workflow(&app_ctx).await;
+
+        let raster_colorizer = RasterColorizer::MultiBand {
+            red_band: 2,
+            red_min: 0.,
+            red_max: 255.,
+            red_scale: 1.0,
+            green_band: 1,
+            green_min: 0.,
+            green_max: 255.,
+            green_scale: 1.0,
+            blue_band: 0,
+            blue_min: 0.,
+            blue_max: 255.,
+            blue_scale: 1.0,
+            no_data_color: RgbaColor::transparent().into(),
+        };
+
+        let params = &[
+            ("request", "GetMap"),
+            ("service", "WMS"),
+            ("version", "1.3.0"),
+            ("layers", &id.to_string()),
+            ("bbox", "-90,-180,90,180"),
+            ("width", "600"),
+            ("height", "300"),
+            ("crs", "EPSG:4326"),
+            (
+                "styles",
+                &format!(
+                    "custom:{}",
+                    serde_json::to_string(&raster_colorizer).unwrap()
+                ),
+            ),
+            ("format", "image/png"),
+            ("time", "2022-01-01T00:00:00.0Z"),
+        ];
+
+        let req = actix_web::test::TestRequest::get()
+            .uri(&format!(
+                "/wms/{}?{}",
+                id,
+                serde_urlencoded::to_string(params).unwrap()
+            ))
+            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));
+        let res = send_test_request(req, app_ctx).await;
+
+        assert_eq!(res.status(), 200);
+
+        let image_bytes = actix_web::test::read_body(res).await;
+
+        // geoengine_datatypes::util::test::save_test_bytes(&image_bytes, "ne2_rgb_colorizer.png");
+
+        assert_eq!(
+            include_bytes!("../../../../test_data/wms/ne2_rgb_colorizer.png") as &[u8],
+            image_bytes
+        );
+    }
+
+    #[ge_context::test(tiling_spec = "get_map_test_helper_tiling_spec")]
+    async fn it_supports_multiband_colorizer_with_less_then_3_bands(
+        app_ctx: PostgresContext<NoTls>,
+    ) {
+        let ctx = app_ctx.default_session_context().await.unwrap();
+
+        let session_id = ctx.session().id();
+
+        let (_, id) = register_ne2_multiband_workflow(&app_ctx).await;
+
+        let raster_colorizer = RasterColorizer::MultiBand {
+            red_band: 1,
+            red_min: 0.,
+            red_max: 255.,
+            red_scale: 1.0,
+            green_band: 1,
+            green_min: 0.,
+            green_max: 255.,
+            green_scale: 1.0,
+            blue_band: 1,
+            blue_min: 0.,
+            blue_max: 255.,
+            blue_scale: 1.0,
+            no_data_color: RgbaColor::transparent().into(),
+        };
+
+        let params = &[
+            ("request", "GetMap"),
+            ("service", "WMS"),
+            ("version", "1.3.0"),
+            ("layers", &id.to_string()),
+            ("bbox", "-90,-180,90,180"),
+            ("width", "600"),
+            ("height", "300"),
+            ("crs", "EPSG:4326"),
+            (
+                "styles",
+                &format!(
+                    "custom:{}",
+                    serde_json::to_string(&raster_colorizer).unwrap()
+                ),
+            ),
+            ("format", "image/png"),
+            ("time", "2022-01-01T00:00:00.0Z"),
+        ];
+
+        let req = actix_web::test::TestRequest::get()
+            .uri(&format!(
+                "/wms/{}?{}",
+                id,
+                serde_urlencoded::to_string(params).unwrap()
+            ))
+            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));
+        let res = send_test_request(req, app_ctx).await;
+
+        assert_eq!(res.status(), 200);
+
+        let image_bytes = actix_web::test::read_body(res).await;
+
+        // geoengine_datatypes::util::test::save_test_bytes(
+        //     &image_bytes,
+        //     geoengine_datatypes::test_data!("wms/ne2_rgb_colorizer_gray.png")
+        //         .to_str()
+        //         .unwrap(),
+        // );
+
+        assert_eq!(
+            include_bytes!("../../../../test_data/wms/ne2_rgb_colorizer_gray.png") as &[u8],
+            image_bytes
+        );
+    }
+
     #[ge_context::test]
     async fn it_zoomes_very_far(app_ctx: PostgresContext<NoTls>) {
         let ctx = app_ctx.default_session_context().await.unwrap();