40592098596

Committed 27 Oct 2026 10:05PM UTC coverage: 87.618% (-0.01%) from 87.631%

Build # 40592098596

Build Type

Pull #1370

github

Committed by

web-flow

Commit Message

Merge 9a5c703f9 into 6ce98f3e4

Pull Request Pull Request #1370: fix: expression should not skip all empty if map_no_data

Coverage Stats

26 of 50 new or added lines in 1 file covered. (52.0%)

3 existing lines in 1 file now uncovered.

134466 of 153468 relevant lines covered (87.62%)

546634.96 hits per line

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

80.36

/geoengine/operators/src/processing/expression/raster_query_processor.rs

use super::RasterExpressionError;
use crate::{
    engine::{
        BoxRasterQueryProcessor, QueryContext, QueryProcessor, RasterQueryProcessor,
        RasterResultDescriptor,
    },
    util::Result,
};
use async_trait::async_trait;
use futures::{StreamExt, TryStreamExt, stream::BoxStream};
use geoengine_datatypes::{
    primitives::{BandSelection, CacheHint, RasterQueryRectangle, TimeInterval},
    raster::{
        ConvertDataType, FromIndexFnParallel, GeoTransform, GridBoundingBox2D, GridIdx2D,
        GridIndexAccess, GridOrEmpty, GridOrEmpty2D, GridShape2D, GridShapeAccess,
        MapElementsParallel, MaskedGrid2D, Pixel, RasterTile2D,
    },
};
use geoengine_expression::LinkedExpression;
use libloading::Symbol;
use num_traits::AsPrimitive;
use std::{marker::PhantomData, sync::Arc};

pub struct ExpressionInput<const N: usize> {
    pub raster: BoxRasterQueryProcessor<f64>,
}

pub struct ExpressionQueryProcessor<TO, Sources>
where
    TO: Pixel,
{
    pub sources: Sources,
    pub result_descriptor: RasterResultDescriptor,
    pub phantom_data: PhantomData<TO>,
    pub program: Arc<LinkedExpression>,
    pub map_no_data: bool,
}

impl<TO, Sources> ExpressionQueryProcessor<TO, Sources>
where
    TO: Pixel,
{
    pub fn new(
        program: LinkedExpression,
        sources: Sources,
        result_descriptor: RasterResultDescriptor,
        map_no_data: bool,
    ) -> Self {
        Self {
            sources,
            result_descriptor,
            program: Arc::new(program),
            phantom_data: PhantomData,
            map_no_data,
        }
    }
}

#[async_trait]
impl<TO, Tuple> QueryProcessor for ExpressionQueryProcessor<TO, Tuple>
where
    TO: Pixel,
    Tuple: ExpressionTupleProcessor<TO>,
{
    type Output = RasterTile2D<TO>;
    type SpatialBounds = GridBoundingBox2D;
    type Selection = BandSelection;
    type ResultDescription = RasterResultDescriptor;

    async fn _query<'b>(
        &'b self,
        query: RasterQueryRectangle,
        ctx: &'b dyn QueryContext,
    ) -> Result<BoxStream<'b, Result<Self::Output>>> {
        // rewrite query to request all input bands from the source. They are all combined in the single output band by means of the expression.
        let source_query = query.select_attributes(BandSelection::first_n(Tuple::num_bands()));

        let stream =
            self.sources
                .zip_bands(source_query, ctx)
                .await?
                .and_then(move |rasters| async move {
                    if !self.map_no_data && Tuple::all_empty(&rasters) {
                        return Ok(Tuple::empty_raster(&rasters));
                    }

                    let (
                        out_time,
                        out_tile_position,
                        out_global_geo_transform,
                        _output_grid_shape,
                        cache_hint,
                    ) = Tuple::metadata(&rasters);

                    let program = self.program.clone();
                    let map_no_data: bool = self.map_no_data;

                    let out = crate::util::spawn_blocking_with_thread_pool(
                        ctx.thread_pool().clone(),
                        move || Tuple::compute_expression(rasters, &program, map_no_data),
                    )
                    .await??;

                    Ok(RasterTile2D::new(
                        out_time,
                        out_tile_position,
                        0,
                        out_global_geo_transform,
                        out,
                        cache_hint,
                    ))
                });

        Ok(stream.boxed())
    }

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

#[async_trait]
impl<TO, Tuple> RasterQueryProcessor for ExpressionQueryProcessor<TO, Tuple>
where
    TO: Pixel,
    Tuple: ExpressionTupleProcessor<TO>,
{
    type RasterType = TO;

    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.sources.time_query(query, ctx).await
    }
}

#[async_trait]
trait ExpressionTupleProcessor<TO: Pixel>: Send + Sync {
    type Tuple: Send + 'static;

    async fn zip_bands<'a>(
        &'a self,
        query: RasterQueryRectangle,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<Self::Tuple>>>;

    fn all_empty(tuple: &Self::Tuple) -> bool;

    fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO>;

    fn metadata(
        tuple: &Self::Tuple,
    ) -> (
        TimeInterval,
        GridIdx2D,
        GeoTransform,
        GridShape2D,
        CacheHint,
    );

    fn compute_expression(
        tuple: Self::Tuple,
        program: &LinkedExpression,
        map_no_data: bool,
    ) -> Result<GridOrEmpty2D<TO>>;

    fn num_bands() -> u32;

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

#[async_trait]
impl<TO> ExpressionTupleProcessor<TO> for ExpressionInput<1>
where
    TO: Pixel,
    f64: AsPrimitive<TO>,
{
    type Tuple = RasterTile2D<f64>;

    #[inline]
    async fn zip_bands<'a>(
        &'a self,
        query: RasterQueryRectangle,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<Self::Tuple>>> {
        self.raster.query(query, ctx).await
    }

    #[inline]
    fn all_empty(tuple: &Self::Tuple) -> bool {
        tuple.grid_array.is_empty()
    }

    #[inline]
    fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO> {
        tuple.clone().convert_data_type()
    }

    #[inline]
    fn metadata(
        tuple: &Self::Tuple,
    ) -> (
        TimeInterval,
        GridIdx2D,
        GeoTransform,
        GridShape2D,
        CacheHint,
    ) {
        let raster = &tuple;

        (
            raster.time,
            raster.tile_position,
            raster.global_geo_transform,
            raster.grid_shape(),
            raster.cache_hint,
        )
    }

    #[inline]
    fn compute_expression(
        raster: Self::Tuple,
        program: &LinkedExpression,
        map_no_data: bool,
    ) -> Result<GridOrEmpty2D<TO>> {
        let expression = unsafe {
            // we have to "trust" that the function has the signature we expect
            program
                .function_1::<Option<f64>>()
                .map_err(RasterExpressionError::from)?
        };

        if raster.is_empty() && !map_no_data {
            return Ok(GridOrEmpty::new_empty_shape(raster.grid_shape()));
        }

        let empty_result = if map_no_data {
            expression(None).map(TO::from_)
        } else {
            None
        };

        if raster.is_empty() {
            return match empty_result {
                Some(value) => Ok(GridOrEmpty::new_grid(MaskedGrid2D::new_filled(
                    raster.grid_shape(),
                    value,
                ))),
                None => Ok(GridOrEmpty::new_empty_shape(raster.grid_shape())),
            };
        }

        let map_fn = |in_value: Option<f64>| {
            // TODO: could be a |in_value: T1| if map no data is false!
            if in_value.is_none() {
                return empty_result;
            }

            let result = expression(in_value);

            result.map(TO::from_)
        };

        let res = raster.grid_array.map_elements_parallel(map_fn);

        Result::Ok(res)
    }

    fn num_bands() -> u32 {
        1
    }

    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.raster.time_query(query, ctx).await
    }
}

// TODO: implement this via macro for 2-8 sources
#[async_trait]
impl<TO> ExpressionTupleProcessor<TO> for ExpressionInput<2>
where
    TO: Pixel,
    f64: AsPrimitive<TO>,
{
    type Tuple = (RasterTile2D<f64>, RasterTile2D<f64>);

    #[inline]
    async fn zip_bands<'a>(
        &'a self,
        query: RasterQueryRectangle,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<Self::Tuple>>> {
        // chunk up the stream to get all bands for a spatial tile at once
        let stream = self.raster.query(query, ctx).await?.chunks(2).map(|chunk| {
            if chunk.len() != 2 {
                // if there are not exactly two tiles, it should mean the last tile was an error and the chunker ended prematurely
                if let Some(Err(e)) = chunk.into_iter().next_back() {
                    return Err(e);
                }
                // if there is no error, the source did not produce all bands, which likely means a bug in an operator
                return Err(crate::error::Error::MustNotHappen {
                    message: "source did not produce all bands (N: 2)".to_string(),
                });
            }

            let [a, b]: [Result<RasterTile2D<f64>>; 2] = chunk
                .try_into()
                .expect("all chunks should be of length 2 because it was checked above");
            match (a, b) {
                (Ok(t0), Ok(t1)) => Ok((t0, t1)),
                (Err(e), _) | (_, Err(e)) => Err(e),
            }
        });

        Ok(stream.boxed())
    }

    #[inline]
    fn all_empty(tuple: &Self::Tuple) -> bool {
        tuple.0.grid_array.is_empty() && tuple.1.grid_array.is_empty()
    }

    #[inline]
    fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO> {
        tuple.0.clone().convert_data_type()
    }

    #[inline]
    fn metadata(
        tuple: &Self::Tuple,
    ) -> (
        TimeInterval,
        GridIdx2D,
        GeoTransform,
        GridShape2D,
        CacheHint,
    ) {
        let raster = &tuple.0;

        (
            raster.time,
            raster.tile_position,
            raster.global_geo_transform,
            raster.grid_shape(),
            tuple.0.cache_hint.merged(&tuple.1.cache_hint),
        )
    }

    #[inline]
    fn compute_expression(
        rasters: Self::Tuple,
        program: &LinkedExpression,
        map_no_data: bool,
    ) -> Result<GridOrEmpty2D<TO>> {
        let expression = unsafe {
            // we have to "trust" that the function has the signature we expect
            program
                .function_2::<Option<f64>, Option<f64>>()
                .map_err(RasterExpressionError::from)?
        };

        if !map_no_data && rasters.0.is_empty() && rasters.1.is_empty() {
            return Ok(GridOrEmpty::new_empty_shape(rasters.0.grid_shape()));
        }

        let empty_result = if map_no_data {
            expression(None, None).map(TO::from_)
        } else {
            None
        };

        if rasters.0.is_empty() && rasters.1.is_empty() {
            return match empty_result {
                Some(value) => Ok(GridOrEmpty::new_grid(MaskedGrid2D::new_filled(
                    rasters.0.grid_shape(),
                    value,
                ))),
                None => Ok(GridOrEmpty::new_empty_shape(rasters.0.grid_shape())),
            };
        }

        let map_fn = |lin_idx: usize| {
            let t0_value = rasters.0.get_at_grid_index_unchecked(lin_idx);
            let t1_value = rasters.1.get_at_grid_index_unchecked(lin_idx);

            if !map_no_data && (t0_value.is_none() || t1_value.is_none()) {
                return None;
            }

            let result = expression(t0_value, t1_value);

            result.map(TO::from_)
        };

        let grid_shape = rasters.0.grid_shape();
        let out = GridOrEmpty::from_index_fn_parallel(&grid_shape, map_fn);

        let out_or_empty = match out {
            GridOrEmpty::Empty(t) => GridOrEmpty::Empty(t),
            GridOrEmpty::Grid(t) => {
                if t.validity_mask.data.iter().all(|&t| !t) {
                    tracing::trace!("Converting empty expression output to EmptyGrid");
                    GridOrEmpty::new_empty_shape(t.grid_shape())
                } else {
                    GridOrEmpty::Grid(t)
                }
            }
        };

        Result::Ok(out_or_empty)
    }

    fn num_bands() -> u32 {
        2
    }

    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.raster.time_query(query, ctx).await
    }
}

type Function3 = fn(Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
type Function4 = fn(Option<f64>, Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
type Function5 = fn(Option<f64>, Option<f64>, Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
type Function6 =
    fn(Option<f64>, Option<f64>, Option<f64>, Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
type Function7 = fn(
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
) -> Option<f64>;
type Function8 = fn(
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
    Option<f64>,
) -> Option<f64>;

macro_rules! impl_expression_tuple_processor {
    ( $i:tt => $( $x:tt ),+ ) => {
        paste::paste! {
            impl_expression_tuple_processor!(
                @inner
                $i
                |
                $( $x ),*
                |
                $( [< pixel_ $x >] ),*
                |
                $( [< is_nodata_ $x >] ),*
                |
                [< Function $i >]
            );
        }
    };

    // We have `0, 1, 2, …` and `T0, T1, T2, …`
    (@inner $N:tt | $( $I:tt ),+ | $( $PIXEL:tt ),+ | $( $IS_NODATA:tt ),+ | $FN_T:ty ) => {
        #[async_trait]
        impl<TO> ExpressionTupleProcessor<TO> for ExpressionInput<$N>
        where
            TO: Pixel,
            f64: AsPrimitive<TO>,
        {
            type Tuple = [RasterTile2D<f64>; $N];

            #[inline]
            async fn zip_bands<'a>(
                &'a self,
                query: RasterQueryRectangle,
                ctx: &'a dyn QueryContext,
            ) -> Result<BoxStream<'a, Result<Self::Tuple>>> {
                // chunk up the stream to get all bands for a spatial tile at once
                let stream = self.raster.query(query, ctx).await?.chunks($N).map(|chunk| {
                    if chunk.len() != $N {
                        // if there are not exactly N tiles, it should mean the last tile was an error and the chunker ended prematurely
                        if let Some(Err(e)) = chunk.into_iter().last() {
                            return Err(e);
                        }
                        // if there is no error, the source did not produce all bands, which likely means a bug in an operator
                        return Err(crate::error::Error::MustNotHappen {
                            message: format!("source did not produce all bands, (N: {})", $N),
                        });
                    }

                    let mut ok_tiles = Vec::with_capacity($N);

                    for tile in chunk {
                        match tile {
                            Ok(tile) => ok_tiles.push(tile),
                            Err(e) => return Err(e),
                        }
                    }

                    let tuple: [RasterTile2D<f64>; $N] = ok_tiles
                        .try_into()
                        .expect("all chunks should be of the expected langth because it was checked above");

                    Ok(tuple)
                });

                Ok(stream.boxed())
            }

            #[inline]
            fn all_empty(tuple: &Self::Tuple) -> bool {
                $( tuple[$I].grid_array.is_empty() )&&*
            }

            #[inline]
            fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO> {
                tuple[0].clone().convert_data_type()
            }

            #[inline]
            fn metadata(tuple: &Self::Tuple) -> (TimeInterval, GridIdx2D, GeoTransform, GridShape2D, CacheHint) {
                let raster = &tuple[0];

                (
                    raster.time,
                    raster.tile_position,
                    raster.global_geo_transform,
                    raster.grid_shape(),
                    tuple.iter().fold(CacheHint::max_duration(), |acc, r| acc.merged(&r.cache_hint)),
                )
            }

            fn compute_expression(
                rasters: Self::Tuple,
                program: &LinkedExpression,
                map_no_data: bool,
            ) -> Result<GridOrEmpty2D<TO>> {
                let expression: Symbol<$FN_T> = unsafe {
                    // we have to "trust" that the function has the signature we expect
                    program.function_nary().map_err(RasterExpressionError::from)?
                };

                let all_empty = rasters.iter().all(|r| r.grid_array.is_empty());

                if all_empty && !map_no_data {
                    return Ok(GridOrEmpty::new_empty_shape(rasters[0].grid_shape()));
                };

                // call expression with $N None values correctly:
                let empty_result = if map_no_data {
                    expression(
                        $(
                            { let _ = $I; None } // this is a nice trick to get the correct number of None values for the function call
                        ),+
                    ).map(TO::from_)
                } else {
                    None
                };

                if all_empty {
                    return match empty_result {
                        Some(value) => Ok(GridOrEmpty::new_grid(MaskedGrid2D::new_filled(
                            rasters[0].grid_shape(),
                            value,
                        ))),
                        None => Ok(GridOrEmpty::new_empty_shape(rasters[0].grid_shape())),
                    };
                }

                let map_fn = |lin_idx: usize| {
                    $(
                        let $PIXEL = rasters[$I].get_at_grid_index_unchecked(lin_idx);
                        let $IS_NODATA = $PIXEL.is_none();
                    )*

                    if !map_no_data && ( $($IS_NODATA)||* ) {
                        return None;
                    }

                    let result = expression(
                        $(
                            $PIXEL
                        ),*
                    );

                    result.map(TO::from_)
                };

                let grid_shape = rasters[0].grid_shape();
                let out = GridOrEmpty::from_index_fn_parallel(&grid_shape, map_fn);

                Result::Ok(out)
            }

            fn num_bands() -> u32 {
                $N
            }

            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.raster.time_query(query, ctx).await
                }
        }


    };

    // For any input, generate `f64, bool`
    (@input_dtypes $x:tt) => {
        f64, bool
    };
}

impl_expression_tuple_processor!(3 => 0, 1, 2);
impl_expression_tuple_processor!(4 => 0, 1, 2, 3);
impl_expression_tuple_processor!(5 => 0, 1, 2, 3, 4);
impl_expression_tuple_processor!(6 => 0, 1, 2, 3, 4, 5);
impl_expression_tuple_processor!(7 => 0, 1, 2, 3, 4, 5, 6);
impl_expression_tuple_processor!(8 => 0, 1, 2, 3, 4, 5, 6, 7);

1	use super::RasterExpressionError;
2	use crate::{
3	engine::{
4	BoxRasterQueryProcessor, QueryContext, QueryProcessor, RasterQueryProcessor,
5	RasterResultDescriptor,
6	},
7	util::Result,
8	};
9	use async_trait::async_trait;
10	use futures::{StreamExt, TryStreamExt, stream::BoxStream};
11	use geoengine_datatypes::{
12	primitives::{BandSelection, CacheHint, RasterQueryRectangle, TimeInterval},
13	raster::{
14	ConvertDataType, FromIndexFnParallel, GeoTransform, GridBoundingBox2D, GridIdx2D,
15	GridIndexAccess, GridOrEmpty, GridOrEmpty2D, GridShape2D, GridShapeAccess,
16	MapElementsParallel, MaskedGrid2D, Pixel, RasterTile2D,
17	},
18	};
19	use geoengine_expression::LinkedExpression;
20	use libloading::Symbol;
21	use num_traits::AsPrimitive;
22	use std::{marker::PhantomData, sync::Arc};
23
24	pub struct ExpressionInput<const N: usize> {
25	pub raster: BoxRasterQueryProcessor<f64>,
26	}
27
28	pub struct ExpressionQueryProcessor<TO, Sources>
29	where
30	TO: Pixel,
31	{
32	pub sources: Sources,
33	pub result_descriptor: RasterResultDescriptor,
34	pub phantom_data: PhantomData<TO>,
35	pub program: Arc<LinkedExpression>,
36	pub map_no_data: bool,
37	}
38
39	impl<TO, Sources> ExpressionQueryProcessor<TO, Sources>
40	where
41	TO: Pixel,
42	{
43	pub fn new(	16 ✔
44	program: LinkedExpression,	16 ✔
45	sources: Sources,	16 ✔
46	result_descriptor: RasterResultDescriptor,	16 ✔
47	map_no_data: bool,	16 ✔
48	) -> Self {	16 ✔
49	Self {	16 ✔
50	sources,	16 ✔
51	result_descriptor,	16 ✔
52	program: Arc::new(program),	16 ✔
53	phantom_data: PhantomData,	16 ✔
54	map_no_data,	16 ✔
55	}	16 ✔
56	}	16 ✔
57	}
58
59	#[async_trait]
60	impl<TO, Tuple> QueryProcessor for ExpressionQueryProcessor<TO, Tuple>
61	where
62	TO: Pixel,
63	Tuple: ExpressionTupleProcessor<TO>,
64	{
65	type Output = RasterTile2D<TO>;
66	type SpatialBounds = GridBoundingBox2D;
67	type Selection = BandSelection;
68	type ResultDescription = RasterResultDescriptor;
69
70	async fn _query<'b>(
71	&'b self,
72	query: RasterQueryRectangle,
73	ctx: &'b dyn QueryContext,
74	) -> Result<BoxStream<'b, Result<Self::Output>>> {	21 ✔
75	// rewrite query to request all input bands from the source. They are all combined in the single output band by means of the expression.
76	let source_query = query.select_attributes(BandSelection::first_n(Tuple::num_bands()));
77
78	let stream =
79	self.sources
80	.zip_bands(source_query, ctx)
81	.await?
82	.and_then(move \|rasters\| async move {	171 ✔
83	if !self.map_no_data && Tuple::all_empty(&rasters) {	171 ✔
84	return Ok(Tuple::empty_raster(&rasters));	16 ✔
85	}	155 ✔
86
87	let (
88	out_time,	155 ✔
89	out_tile_position,	155 ✔
90	out_global_geo_transform,	155 ✔
91	_output_grid_shape,	155 ✔
92	cache_hint,	155 ✔
93	) = Tuple::metadata(&rasters);	155 ✔
94
95	let program = self.program.clone();	155 ✔
96	let map_no_data: bool = self.map_no_data;	155 ✔
97
98	let out = crate::util::spawn_blocking_with_thread_pool(	155 ✔
99	ctx.thread_pool().clone(),	155 ✔
100	move \|\| Tuple::compute_expression(rasters, &program, map_no_data),	155 ✔
101	)
102	.await??;	155 ✔
103
104	Ok(RasterTile2D::new(	155 ✔
105	out_time,	155 ✔
106	out_tile_position,	155 ✔
107	0,	155 ✔
108	out_global_geo_transform,	155 ✔
109	out,	155 ✔
110	cache_hint,	155 ✔
111	))	155 ✔
112	});	342 ✔
113
114	Ok(stream.boxed())
115	}	21 ✔
116
117	fn result_descriptor(&self) -> &RasterResultDescriptor {	63 ✔
118	&self.result_descriptor	63 ✔
119	}	63 ✔
120	}
121
122	#[async_trait]
123	impl<TO, Tuple> RasterQueryProcessor for ExpressionQueryProcessor<TO, Tuple>
124	where
125	TO: Pixel,
126	Tuple: ExpressionTupleProcessor<TO>,
127	{
128	type RasterType = TO;
129
130	async fn _time_query<'a>(
131	&'a self,
132	query: geoengine_datatypes::primitives::TimeInterval,
133	ctx: &'a dyn QueryContext,
134	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {	×
135	self.sources.time_query(query, ctx).await
136	}	×
137	}
138
139	#[async_trait]
140	trait ExpressionTupleProcessor<TO: Pixel>: Send + Sync {
141	type Tuple: Send + 'static;
142
143	async fn zip_bands<'a>(
144	&'a self,
145	query: RasterQueryRectangle,
146	ctx: &'a dyn QueryContext,
147	) -> Result<BoxStream<'a, Result<Self::Tuple>>>;
148
149	fn all_empty(tuple: &Self::Tuple) -> bool;
150
151	fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO>;
152
153	fn metadata(
154	tuple: &Self::Tuple,
155	) -> (
156	TimeInterval,
157	GridIdx2D,
158	GeoTransform,
159	GridShape2D,
160	CacheHint,
161	);
162
163	fn compute_expression(
164	tuple: Self::Tuple,
165	program: &LinkedExpression,
166	map_no_data: bool,
167	) -> Result<GridOrEmpty2D<TO>>;
168
169	fn num_bands() -> u32;
170
171	async fn time_query<'a>(
172	&'a self,
173	query: geoengine_datatypes::primitives::TimeInterval,
174	ctx: &'a dyn QueryContext,
175	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>>;
176	}
177
178	#[async_trait]
179	impl<TO> ExpressionTupleProcessor<TO> for ExpressionInput<1>
180	where
181	TO: Pixel,
182	f64: AsPrimitive<TO>,
183	{
184	type Tuple = RasterTile2D<f64>;
185
186	#[inline]
187	async fn zip_bands<'a>(
188	&'a self,
189	query: RasterQueryRectangle,
190	ctx: &'a dyn QueryContext,
191	) -> Result<BoxStream<'a, Result<Self::Tuple>>> {	13 ✔
192	self.raster.query(query, ctx).await
193	}	13 ✔
194
195	#[inline]
196	fn all_empty(tuple: &Self::Tuple) -> bool {	132 ✔
197	tuple.grid_array.is_empty()	132 ✔
198	}	132 ✔
199
200	#[inline]
201	fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO> {	16 ✔
202	tuple.clone().convert_data_type()	16 ✔
203	}	16 ✔
204
205	#[inline]
206	fn metadata(	125 ✔
207	tuple: &Self::Tuple,	125 ✔
208	) -> (	125 ✔
209	TimeInterval,	125 ✔
210	GridIdx2D,	125 ✔
211	GeoTransform,	125 ✔
212	GridShape2D,	125 ✔
213	CacheHint,	125 ✔
214	) {	125 ✔
215	let raster = &tuple;	125 ✔
216
217	(	125 ✔
218	raster.time,	125 ✔
219	raster.tile_position,	125 ✔
220	raster.global_geo_transform,	125 ✔
221	raster.grid_shape(),	125 ✔
222	raster.cache_hint,	125 ✔
223	)	125 ✔
224	}	125 ✔
225
226	#[inline]
227	fn compute_expression(	125 ✔
228	raster: Self::Tuple,	125 ✔
229	program: &LinkedExpression,	125 ✔
230	map_no_data: bool,	125 ✔
231	) -> Result<GridOrEmpty2D<TO>> {	125 ✔
232	let expression = unsafe {	125 ✔
233	// we have to "trust" that the function has the signature we expect
234	program	125 ✔
235	.function_1::<Option<f64>>()	125 ✔
236	.map_err(RasterExpressionError::from)?	125 ✔
237	};
238
239	if raster.is_empty() && !map_no_data {	125 ✔
NEW 240	return Ok(GridOrEmpty::new_empty_shape(raster.grid_shape()));	×
241	}	125 ✔
242
243	let empty_result = if map_no_data {	125 ✔
244	expression(None).map(TO::from_)	9 ✔
245	} else {
246	None	116 ✔
247	};
248
249	if raster.is_empty() {	125 ✔
NEW 250	return match empty_result {	×
NEW 251	Some(value) => Ok(GridOrEmpty::new_grid(MaskedGrid2D::new_filled(	×
NEW 252	raster.grid_shape(),	×
NEW 253	value,	×
NEW 254	))),	×
NEW 255	None => Ok(GridOrEmpty::new_empty_shape(raster.grid_shape())),	×
256	};
257	}	125 ✔
258
259	let map_fn = \|in_value: Option<f64>\| {	29,360,206 ✔
260	// TODO: could be a \|in_value: T1\| if map no data is false!
261	if in_value.is_none() {	29,360,206 ✔
262	return empty_result;	3,859,188 ✔
263	}	25,501,018 ✔
264
265	let result = expression(in_value);	25,501,018 ✔
266
267	result.map(TO::from_)	25,501,018 ✔
268	};	29,360,206 ✔
269
270	let res = raster.grid_array.map_elements_parallel(map_fn);	125 ✔
271
272	Result::Ok(res)	125 ✔
273	}	125 ✔
274
275	fn num_bands() -> u32 {	13 ✔
276	1	13 ✔
277	}	13 ✔
278
279	async fn time_query<'a>(
280	&'a self,
281	query: geoengine_datatypes::primitives::TimeInterval,
282	ctx: &'a dyn QueryContext,
283	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {	×
284	self.raster.time_query(query, ctx).await
285	}	×
286	}
287
288	// TODO: implement this via macro for 2-8 sources
289	#[async_trait]
290	impl<TO> ExpressionTupleProcessor<TO> for ExpressionInput<2>
291	where
292	TO: Pixel,
293	f64: AsPrimitive<TO>,
294	{
295	type Tuple = (RasterTile2D<f64>, RasterTile2D<f64>);
296
297	#[inline]
298	async fn zip_bands<'a>(
299	&'a self,
300	query: RasterQueryRectangle,
301	ctx: &'a dyn QueryContext,
302	) -> Result<BoxStream<'a, Result<Self::Tuple>>> {	5 ✔
303	// chunk up the stream to get all bands for a spatial tile at once
304	let stream = self.raster.query(query, ctx).await?.chunks(2).map(\|chunk\| {	27 ✔
305	if chunk.len() != 2 {	27 ✔
306	// if there are not exactly two tiles, it should mean the last tile was an error and the chunker ended prematurely
307	if let Some(Err(e)) = chunk.into_iter().next_back() {	×
308	return Err(e);	×
309	}	×
310	// if there is no error, the source did not produce all bands, which likely means a bug in an operator
311	return Err(crate::error::Error::MustNotHappen {	×
312	message: "source did not produce all bands (N: 2)".to_string(),	×
313	});	×
314	}	27 ✔
315
316	let [a, b]: [Result<RasterTile2D<f64>>; 2] = chunk	27 ✔
317	.try_into()	27 ✔
318	.expect("all chunks should be of length 2 because it was checked above");	27 ✔
319	match (a, b) {	27 ✔
320	(Ok(t0), Ok(t1)) => Ok((t0, t1)),	27 ✔
321	(Err(e), _) \| (_, Err(e)) => Err(e),	×
322	}
323	});	27 ✔
324
325	Ok(stream.boxed())
326	}	5 ✔
327
328	#[inline]
329	fn all_empty(tuple: &Self::Tuple) -> bool {	6 ✔
330	tuple.0.grid_array.is_empty() && tuple.1.grid_array.is_empty()	6 ✔
331	}	6 ✔
332
333	#[inline]
UNCOV 334	fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO> {	×
UNCOV 335	tuple.0.clone().convert_data_type()	×
UNCOV 336	}	×
337
338	#[inline]
339	fn metadata(	27 ✔
340	tuple: &Self::Tuple,	27 ✔
341	) -> (	27 ✔
342	TimeInterval,	27 ✔
343	GridIdx2D,	27 ✔
344	GeoTransform,	27 ✔
345	GridShape2D,	27 ✔
346	CacheHint,	27 ✔
347	) {	27 ✔
348	let raster = &tuple.0;	27 ✔
349
350	(	27 ✔
351	raster.time,	27 ✔
352	raster.tile_position,	27 ✔
353	raster.global_geo_transform,	27 ✔
354	raster.grid_shape(),	27 ✔
355	tuple.0.cache_hint.merged(&tuple.1.cache_hint),	27 ✔
356	)	27 ✔
357	}	27 ✔
358
359	#[inline]
360	fn compute_expression(	27 ✔
361	rasters: Self::Tuple,	27 ✔
362	program: &LinkedExpression,	27 ✔
363	map_no_data: bool,	27 ✔
364	) -> Result<GridOrEmpty2D<TO>> {	27 ✔
365	let expression = unsafe {	27 ✔
366	// we have to "trust" that the function has the signature we expect
367	program	27 ✔
368	.function_2::<Option<f64>, Option<f64>>()	27 ✔
369	.map_err(RasterExpressionError::from)?	27 ✔
370	};
371
372	if !map_no_data && rasters.0.is_empty() && rasters.1.is_empty() {	27 ✔
NEW 373	return Ok(GridOrEmpty::new_empty_shape(rasters.0.grid_shape()));	×
374	}	27 ✔
375
376	let empty_result = if map_no_data {	27 ✔
377	expression(None, None).map(TO::from_)	21 ✔
378	} else {
379	None	6 ✔
380	};
381
382	if rasters.0.is_empty() && rasters.1.is_empty() {	27 ✔
383	return match empty_result {	12 ✔
NEW 384	Some(value) => Ok(GridOrEmpty::new_grid(MaskedGrid2D::new_filled(	×
NEW 385	rasters.0.grid_shape(),	×
NEW 386	value,	×
NEW 387	))),	×
388	None => Ok(GridOrEmpty::new_empty_shape(rasters.0.grid_shape())),	12 ✔
389	};
390	}	15 ✔
391
392	let map_fn = \|lin_idx: usize\| {	3,145,746 ✔
393	let t0_value = rasters.0.get_at_grid_index_unchecked(lin_idx);	3,145,746 ✔
394	let t1_value = rasters.1.get_at_grid_index_unchecked(lin_idx);	3,145,746 ✔
395
396	if !map_no_data && (t0_value.is_none() \|\| t1_value.is_none()) {	3,145,746 ✔
397	return None;	99 ✔
398	}	3,145,647 ✔
399
400	let result = expression(t0_value, t1_value);	3,145,647 ✔
401
402	result.map(TO::from_)	3,145,647 ✔
403	};	3,145,746 ✔
404
405	let grid_shape = rasters.0.grid_shape();	15 ✔
406	let out = GridOrEmpty::from_index_fn_parallel(&grid_shape, map_fn);	15 ✔
407
408	let out_or_empty = match out {	15 ✔
409	GridOrEmpty::Empty(t) => GridOrEmpty::Empty(t),	×
410	GridOrEmpty::Grid(t) => {	15 ✔
411	if t.validity_mask.data.iter().all(\|&t\| !t) {	972,767 ✔
412	tracing::trace!("Converting empty expression output to EmptyGrid");	×
413	GridOrEmpty::new_empty_shape(t.grid_shape())	×
414	} else {
415	GridOrEmpty::Grid(t)	15 ✔
416	}
417	}
418	};
419
420	Result::Ok(out_or_empty)	15 ✔
421	}	27 ✔
422
423	fn num_bands() -> u32 {	5 ✔
424	2	5 ✔
425	}	5 ✔
426
427	async fn time_query<'a>(
428	&'a self,
429	query: geoengine_datatypes::primitives::TimeInterval,
430	ctx: &'a dyn QueryContext,
431	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {	×
432	self.raster.time_query(query, ctx).await
433	}	×
434	}
435
436	type Function3 = fn(Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
437	type Function4 = fn(Option<f64>, Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
438	type Function5 = fn(Option<f64>, Option<f64>, Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
439	type Function6 =
440	fn(Option<f64>, Option<f64>, Option<f64>, Option<f64>, Option<f64>, Option<f64>) -> Option<f64>;
441	type Function7 = fn(
442	Option<f64>,
443	Option<f64>,
444	Option<f64>,
445	Option<f64>,
446	Option<f64>,
447	Option<f64>,
448	Option<f64>,
449	) -> Option<f64>;
450	type Function8 = fn(
451	Option<f64>,
452	Option<f64>,
453	Option<f64>,
454	Option<f64>,
455	Option<f64>,
456	Option<f64>,
457	Option<f64>,
458	Option<f64>,
459	) -> Option<f64>;
460
461	macro_rules! impl_expression_tuple_processor {
462	( $i:tt => $( $x:tt ),+ ) => {
463	paste::paste! {
464	impl_expression_tuple_processor!(
465	@inner
466	$i
467	\|
468	$( $x ),*
469	\|
470	$( [< pixel_ $x >] ),*
471	\|
472	$( [< is_nodata_ $x >] ),*
473	\|
474	[< Function $i >]
475	);
476	}
477	};
478
479	// We have `0, 1, 2, …` and `T0, T1, T2, …`
480	(@inner $N:tt \| $( $I:tt ),+ \| $( $PIXEL:tt ),+ \| $( $IS_NODATA:tt ),+ \| $FN_T:ty ) => {
481	#[async_trait]
482	impl<TO> ExpressionTupleProcessor<TO> for ExpressionInput<$N>
483	where
484	TO: Pixel,
485	f64: AsPrimitive<TO>,
486	{
487	type Tuple = [RasterTile2D<f64>; $N];
488
489	#[inline]
490	async fn zip_bands<'a>(
491	&'a self,
492	query: RasterQueryRectangle,
493	ctx: &'a dyn QueryContext,
494	) -> Result<BoxStream<'a, Result<Self::Tuple>>> {
495	// chunk up the stream to get all bands for a spatial tile at once
496	let stream = self.raster.query(query, ctx).await?.chunks($N).map(\|chunk\| {	3 ✔
497	if chunk.len() != $N {	1 ✔
498	// if there are not exactly N tiles, it should mean the last tile was an error and the chunker ended prematurely
499	if let Some(Err(e)) = chunk.into_iter().last() {	×
500	return Err(e);	×
501	}	×
502	// if there is no error, the source did not produce all bands, which likely means a bug in an operator
503	return Err(crate::error::Error::MustNotHappen {	×
504	message: format!("source did not produce all bands, (N: {})", $N),	×
505	});	×
506	}	3 ✔
507
508	let mut ok_tiles = Vec::with_capacity($N);	3 ✔
509
510	for tile in chunk {	14 ✔
511	match tile {	14 ✔
512	Ok(tile) => ok_tiles.push(tile),	14 ✔
513	Err(e) => return Err(e),	×
514	}
515	}
516
517	let tuple: [RasterTile2D<f64>; $N] = ok_tiles	3 ✔
518	.try_into()	3 ✔
519	.expect("all chunks should be of the expected langth because it was checked above");	3 ✔
520
521	Ok(tuple)	3 ✔
522	});	3 ✔
523
524	Ok(stream.boxed())
525	}	3 ✔
526
527	#[inline]
528	fn all_empty(tuple: &Self::Tuple) -> bool {	3 ✔
529	$( tuple[$I].grid_array.is_empty() )&&*	×
530	}	3 ✔
531
532	#[inline]
533	fn empty_raster(tuple: &Self::Tuple) -> RasterTile2D<TO> {	×
534	tuple[0].clone().convert_data_type()	×
535	}	×
536
537	#[inline]
538	fn metadata(tuple: &Self::Tuple) -> (TimeInterval, GridIdx2D, GeoTransform, GridShape2D, CacheHint) {	3 ✔
539	let raster = &tuple[0];	3 ✔
540
541	(
542	raster.time,	3 ✔
543	raster.tile_position,	3 ✔
544	raster.global_geo_transform,	3 ✔
545	raster.grid_shape(),	3 ✔
546	tuple.iter().fold(CacheHint::max_duration(), \|acc, r\| acc.merged(&r.cache_hint)),	14 ✔
547	)
548	}	3 ✔
549
550	fn compute_expression(	3 ✔
551	rasters: Self::Tuple,	3 ✔
552	program: &LinkedExpression,	3 ✔
553	map_no_data: bool,	3 ✔
554	) -> Result<GridOrEmpty2D<TO>> {	3 ✔
555	let expression: Symbol<$FN_T> = unsafe {	3 ✔
556	// we have to "trust" that the function has the signature we expect
557	program.function_nary().map_err(RasterExpressionError::from)?	3 ✔
558	};
559
560	let all_empty = rasters.iter().all(\|r\| r.grid_array.is_empty());	3 ✔
561
562	if all_empty && !map_no_data {	3 ✔
NEW 563	return Ok(GridOrEmpty::new_empty_shape(rasters[0].grid_shape()));	×
564	};	3 ✔
565
566	// call expression with $N None values correctly:
567	let empty_result = if map_no_data {	3 ✔
NEW 568	expression(	×
NEW 569	$(	×
NEW 570	{ let _ = $I; None } // this is a nice trick to get the correct number of None values for the function call	×
NEW 571	),+	×
NEW 572	).map(TO::from_)	×
573	} else {
574	None	3 ✔
575	};
576
577	if all_empty {	3 ✔
NEW 578	return match empty_result {	×
NEW 579	Some(value) => Ok(GridOrEmpty::new_grid(MaskedGrid2D::new_filled(	×
NEW 580	rasters[0].grid_shape(),	×
NEW 581	value,	×
NEW 582	))),	×
NEW 583	None => Ok(GridOrEmpty::new_empty_shape(rasters[0].grid_shape())),	×
584	};
585	}	3 ✔
586
587	let map_fn = \|lin_idx: usize\| {	18 ✔
588	$(
589	let $PIXEL = rasters[$I].get_at_grid_index_unchecked(lin_idx);	18 ✔
590	let $IS_NODATA = $PIXEL.is_none();	18 ✔
591	)*
592
593	if !map_no_data && ( $($IS_NODATA)\|\|* ) {	18 ✔
594	return None;	1 ✔
595	}	17 ✔
596
597	let result = expression(	17 ✔
598	$(	17 ✔
599	$PIXEL	17 ✔
600	),*	17 ✔
601	);	17 ✔
602
603	result.map(TO::from_)	17 ✔
604	};	18 ✔
605
606	let grid_shape = rasters[0].grid_shape();	3 ✔
607	let out = GridOrEmpty::from_index_fn_parallel(&grid_shape, map_fn);	3 ✔
608
609	Result::Ok(out)	3 ✔
610	}	3 ✔
611
612	fn num_bands() -> u32 {	3 ✔
613	$N
614	}	3 ✔
615
616	async fn time_query<'a>(
617	&'a self,
618	query: geoengine_datatypes::primitives::TimeInterval,
619	ctx: &'a dyn QueryContext,
620	) -> Result<BoxStream<'a, Result<geoengine_datatypes::primitives::TimeInterval>>> {
621	self.raster.time_query(query, ctx).await
622	}	×
623	}
624
625
626	};
627
628	// For any input, generate `f64, bool`
629	(@input_dtypes $x:tt) => {
630	f64, bool
631	};
632	}
633
634	impl_expression_tuple_processor!(3 => 0, 1, 2);
635	impl_expression_tuple_processor!(4 => 0, 1, 2, 3);
636	impl_expression_tuple_processor!(5 => 0, 1, 2, 3, 4);
637	impl_expression_tuple_processor!(6 => 0, 1, 2, 3, 4, 5);
638	impl_expression_tuple_processor!(7 => 0, 1, 2, 3, 4, 5, 6);
639	impl_expression_tuple_processor!(8 => 0, 1, 2, 3, 4, 5, 6, 7);

geo-engine / geoengine / 40592098596

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