12469296660

Committed 23 Dec 2024 03:15PM UTC coverage: 90.56% (-0.1%) from 90.695%

Build # 12469296660

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #998 from geo-engine/quota_log_wip

Quota and Data usage Logging

Run Details

859 of 1214 new or added lines in 66 files covered. (70.76%)

3 existing lines in 2 files now uncovered.

133923 of 147883 relevant lines covered (90.56%)

54439.32 hits per line

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

95.86

/operators/src/processing/time_shift.rs

use crate::engine::{
    CanonicOperatorName, ExecutionContext, InitializedRasterOperator,
    InitializedSingleRasterOrVectorOperator, InitializedSources, InitializedVectorOperator,
    Operator, OperatorName, QueryContext, RasterOperator, RasterQueryProcessor,
    RasterResultDescriptor, ResultDescriptor, SingleRasterOrVectorSource,
    TypedRasterQueryProcessor, TypedVectorQueryProcessor, VectorOperator, VectorQueryProcessor,
    VectorResultDescriptor, WorkflowOperatorPath,
};
use crate::util::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::StreamExt;
use geoengine_datatypes::collections::{
    FeatureCollection, FeatureCollectionInfos, FeatureCollectionModifications,
};
use geoengine_datatypes::error::{BoxedResultExt, ErrorSource};
use geoengine_datatypes::primitives::{
    ColumnSelection, Duration, Geometry, RasterQueryRectangle, TimeGranularity, TimeInstance,
    TimeInterval,
};
use geoengine_datatypes::primitives::{TimeStep, VectorQueryRectangle};
use geoengine_datatypes::raster::{Pixel, RasterTile2D};
use geoengine_datatypes::util::arrow::ArrowTyped;
use serde::{Deserialize, Serialize};
use snafu::Snafu;

/// Project the query rectangle to a new time interval.
pub type TimeShift = Operator<TimeShiftParams, SingleRasterOrVectorSource>;

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

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum TimeShiftParams {
    /// Shift the query rectangle relative with a time step
    #[serde(rename_all = "camelCase")]
    Relative {
        granularity: TimeGranularity,
        value: i32,
    },
    /// Set the time interval to a fixed value
    #[serde(rename_all = "camelCase")]
    Absolute { time_interval: TimeInterval },
}

pub trait TimeShiftOperation: Send + Sync + Copy {
    type State: Send + Sync + Copy;

    fn shift(
        &self,
        time_interval: TimeInterval,
    ) -> Result<(TimeInterval, Self::State), TimeShiftError>;

    fn reverse_shift(
        &self,
        time_interval: TimeInterval,
        state: Self::State,
    ) -> Result<TimeInterval, TimeShiftError>;
}

#[derive(Debug, Clone, Copy)]
pub struct RelativeForwardShift {
    step: TimeStep,
}

impl TimeShiftOperation for RelativeForwardShift {
    type State = ();

    fn shift(
        &self,
        time_interval: TimeInterval,
    ) -> Result<(TimeInterval, Self::State), TimeShiftError> {
        let time_interval = time_interval + self.step;
        let time_interval = time_interval.boxed_context(error::TimeOverflow)?;
        Ok((time_interval, ()))
    }

    fn reverse_shift(
        &self,
        time_interval: TimeInterval,
        _state: Self::State,
    ) -> Result<TimeInterval, TimeShiftError> {
        let reversed_time_interval = time_interval - self.step;
        reversed_time_interval.boxed_context(error::TimeOverflow)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct RelativeBackwardShift {
    step: TimeStep,
}

impl TimeShiftOperation for RelativeBackwardShift {
    type State = ();

    fn shift(
        &self,
        time_interval: TimeInterval,
    ) -> Result<(TimeInterval, Self::State), TimeShiftError> {
        let time_interval = time_interval - self.step;
        let time_interval = time_interval.boxed_context(error::TimeOverflow)?;
        Ok((time_interval, ()))
    }

    fn reverse_shift(
        &self,
        time_interval: TimeInterval,
        _state: Self::State,
    ) -> Result<TimeInterval, TimeShiftError> {
        let reversed_time_interval = time_interval + self.step;
        reversed_time_interval.boxed_context(error::TimeOverflow)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct AbsoluteShift {
    time_interval: TimeInterval,
}

impl TimeShiftOperation for AbsoluteShift {
    type State = (Duration, Duration);

    fn shift(
        &self,
        time_interval: TimeInterval,
    ) -> Result<(TimeInterval, Self::State), TimeShiftError> {
        let time_start_difference = time_interval.start() - self.time_interval.start();
        let time_end_difference = time_interval.end() - self.time_interval.end();

        Ok((
            self.time_interval,
            (time_start_difference, time_end_difference),
        ))
    }

    fn reverse_shift(
        &self,
        time_interval: TimeInterval,
        (time_start_difference, time_end_difference): Self::State,
    ) -> Result<TimeInterval, TimeShiftError> {
        let t1 = time_interval.start() + time_start_difference.num_milliseconds();
        let t2 = time_interval.end() + time_end_difference.num_milliseconds();
        TimeInterval::new(t1, t2).boxed_context(error::FaultyTimeInterval { t1, t2 })
    }
}

#[derive(Debug, Snafu)]
#[snafu(visibility(pub(crate)), context(suffix(false)), module(error))]
pub enum TimeShiftError {
    #[snafu(display("Output type must match the type of the source"))]
    UnmatchedOutput,
    #[snafu(display("Shifting the time led to an overflowing time interval"))]
    TimeOverflow { source: Box<dyn ErrorSource> },
    #[snafu(display("Shifting the time to a faulty time interval: {t1} / {t2}"))]
    FaultyTimeInterval {
        source: Box<dyn ErrorSource>,
        t1: TimeInstance,
        t2: TimeInstance,
    },
    #[snafu(display("Modifying the timestamps of the feature collection failed"))]
    FeatureCollectionTimeModification { source: Box<dyn ErrorSource> },
}

#[typetag::serde]
#[async_trait]
impl VectorOperator for TimeShift {
    async fn _initialize(
        self: Box<Self>,
        path: WorkflowOperatorPath,
        context: &dyn ExecutionContext,
    ) -> Result<Box<dyn InitializedVectorOperator>> {
        let name = CanonicOperatorName::from(&self);

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

        match (init_sources.source, self.params) {
            (
                InitializedSingleRasterOrVectorOperator::Vector(source),
                TimeShiftParams::Relative { granularity, value },
            ) if value.is_positive() => {
                let shift = RelativeForwardShift {
                    step: TimeStep {
                        granularity,
                        step: value.unsigned_abs(),
                    },
                };

                let result_descriptor = shift_result_descriptor(source.result_descriptor(), shift);

                Ok(Box::new(InitializedVectorTimeShift {
                    name,
                    path,
                    source,
                    result_descriptor,
                    shift,
                }))
            }
            (
                InitializedSingleRasterOrVectorOperator::Vector(source),
                TimeShiftParams::Relative { granularity, value },
            ) => {
                let shift = RelativeBackwardShift {
                    step: TimeStep {
                        granularity,
                        step: value.unsigned_abs(),
                    },
                };

                let result_descriptor = shift_result_descriptor(source.result_descriptor(), shift);

                Ok(Box::new(InitializedVectorTimeShift {
                    name,
                    path,
                    source,
                    result_descriptor,
                    shift,
                }))
            }
            (
                InitializedSingleRasterOrVectorOperator::Vector(source),
                TimeShiftParams::Absolute { time_interval },
            ) => {
                let shift = AbsoluteShift { time_interval };

                let result_descriptor = shift_result_descriptor(source.result_descriptor(), shift);

                Ok(Box::new(InitializedVectorTimeShift {
                    name,
                    path,
                    source,
                    result_descriptor,
                    shift,
                }))
            }
            (InitializedSingleRasterOrVectorOperator::Raster(_), _) => {
                Err(TimeShiftError::UnmatchedOutput.into())
            }
        }
    }

    span_fn!(TimeShift);
}

#[typetag::serde]
#[async_trait]
impl RasterOperator for TimeShift {
    async fn _initialize(
        self: Box<Self>,
        path: WorkflowOperatorPath,
        context: &dyn ExecutionContext,
    ) -> Result<Box<dyn InitializedRasterOperator>> {
        let name = CanonicOperatorName::from(&self);

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

        match (init_sources.source, self.params) {
            (
                InitializedSingleRasterOrVectorOperator::Raster(source),
                TimeShiftParams::Relative { granularity, value },
            ) if value.is_positive() => {
                let shift = RelativeForwardShift {
                    step: TimeStep {
                        granularity,
                        step: value.unsigned_abs(),
                    },
                };

                let result_descriptor = shift_result_descriptor(source.result_descriptor(), shift);

                Ok(Box::new(InitializedRasterTimeShift {
                    name,
                    path,
                    source,
                    result_descriptor,
                    shift,
                }))
            }
            (
                InitializedSingleRasterOrVectorOperator::Raster(source),
                TimeShiftParams::Relative { granularity, value },
            ) => {
                let shift = RelativeBackwardShift {
                    step: TimeStep {
                        granularity,
                        step: value.unsigned_abs(),
                    },
                };

                let result_descriptor = shift_result_descriptor(source.result_descriptor(), shift);

                Ok(Box::new(InitializedRasterTimeShift {
                    name,
                    path,
                    source,
                    result_descriptor,
                    shift,
                }))
            }
            (
                InitializedSingleRasterOrVectorOperator::Raster(source),
                TimeShiftParams::Absolute { time_interval },
            ) => {
                let shift = AbsoluteShift { time_interval };

                let result_descriptor = shift_result_descriptor(source.result_descriptor(), shift);

                Ok(Box::new(InitializedRasterTimeShift {
                    name,
                    path,
                    source,
                    result_descriptor,
                    shift,
                }))
            }
            (InitializedSingleRasterOrVectorOperator::Vector(_), _) => {
                Err(TimeShiftError::UnmatchedOutput.into())
            }
        }
    }

    span_fn!(TimeShift);
}

fn shift_result_descriptor<R: ResultDescriptor, S: TimeShiftOperation>(
    result_descriptor: &R,
    shift: S,
) -> R {
    result_descriptor.map_time(|time| {
        if let Some(time) = time {
            shift.shift(*time).map(|r| r.0).ok()
        } else {
            None
        }
    })
}

pub struct InitializedVectorTimeShift<Shift: TimeShiftOperation> {
    name: CanonicOperatorName,
    path: WorkflowOperatorPath,
    source: Box<dyn InitializedVectorOperator>,
    result_descriptor: VectorResultDescriptor,
    shift: Shift,
}

pub struct InitializedRasterTimeShift<Shift: TimeShiftOperation> {
    name: CanonicOperatorName,
    path: WorkflowOperatorPath,
    source: Box<dyn InitializedRasterOperator>,
    result_descriptor: RasterResultDescriptor,
    shift: Shift,
}

impl<Shift: TimeShiftOperation + 'static> InitializedVectorOperator
    for InitializedVectorTimeShift<Shift>
{
    fn result_descriptor(&self) -> &VectorResultDescriptor {
        &self.result_descriptor
    }

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

        Ok(
            call_on_generic_vector_processor!(source_processor, processor => VectorTimeShiftProcessor {
                processor,
                result_descriptor: self.result_descriptor.clone(),
                shift: self.shift,
            }.boxed().into()),
        )
    }

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

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

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

impl<Shift: TimeShiftOperation + 'static> InitializedRasterOperator
    for InitializedRasterTimeShift<Shift>
{
    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, processor => RasterTimeShiftProcessor {
                processor,
                result_descriptor: self.result_descriptor.clone(),
                shift: self.shift,
            }.boxed().into()),
        )
    }

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

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

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

pub struct RasterTimeShiftProcessor<Q, P, Shift: TimeShiftOperation>
where
    Q: RasterQueryProcessor<RasterType = P>,
{
    processor: Q,
    result_descriptor: RasterResultDescriptor,
    shift: Shift,
}

pub struct VectorTimeShiftProcessor<Q, G, Shift: TimeShiftOperation>
where
    G: Geometry,
    Q: VectorQueryProcessor<VectorType = FeatureCollection<G>>,
{
    processor: Q,
    result_descriptor: VectorResultDescriptor,
    shift: Shift,
}

#[async_trait]
impl<Q, G, Shift> VectorQueryProcessor for VectorTimeShiftProcessor<Q, G, Shift>
where
    G: Geometry + ArrowTyped + 'static,
    Q: VectorQueryProcessor<VectorType = FeatureCollection<G>>,
    Shift: TimeShiftOperation + 'static,
{
    type VectorType = FeatureCollection<G>;

    async fn vector_query<'a>(
        &'a self,
        query: VectorQueryRectangle,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<Self::VectorType>>> {
        let (time_interval, state) = self.shift.shift(query.time_interval)?;

        let query = VectorQueryRectangle {
            spatial_bounds: query.spatial_bounds,
            time_interval,
            spatial_resolution: query.spatial_resolution,
            attributes: ColumnSelection::all(),
        };
        let stream = self.processor.vector_query(query, ctx).await?;

        let stream = stream.then(move |collection| async move {
            let collection = collection?;
            let shift = self.shift;

            crate::util::spawn_blocking(move || {
                let time_intervals = collection
                    .time_intervals()
                    .iter()
                    .map(move |time| shift.reverse_shift(*time, state))
                    .collect::<Result<Vec<TimeInterval>, TimeShiftError>>()?;

                collection
                    .replace_time(&time_intervals)
                    .boxed_context(error::FeatureCollectionTimeModification)
                    .map_err(Into::into)
            })
            .await?
        });

        Ok(stream.boxed())
    }

    fn vector_result_descriptor(&self) -> &VectorResultDescriptor {
        &self.result_descriptor
    }
}

#[async_trait]
impl<Q, P, Shift> RasterQueryProcessor for RasterTimeShiftProcessor<Q, P, Shift>
where
    Q: RasterQueryProcessor<RasterType = P>,
    P: Pixel,
    Shift: TimeShiftOperation,
{
    type RasterType = P;

    async fn raster_query<'a>(
        &'a self,
        query: RasterQueryRectangle,
        ctx: &'a dyn QueryContext,
    ) -> Result<BoxStream<'a, Result<RasterTile2D<Self::RasterType>>>> {
        let (time_interval, state) = self.shift.shift(query.time_interval)?;
        let query = RasterQueryRectangle {
            spatial_bounds: query.spatial_bounds,
            time_interval,
            spatial_resolution: query.spatial_resolution,
            attributes: query.attributes,
        };
        let stream = self.processor.raster_query(query, ctx).await?;

        let stream = stream.map(move |raster| {
            // reverse time shift for results
            let mut raster = raster?;

            raster.time = self.shift.reverse_shift(raster.time, state)?;

            Ok(raster)
        });

        Ok(Box::pin(stream))
    }

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

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

    use crate::{
        engine::{
            MockExecutionContext, MockQueryContext, MultipleRasterSources, RasterBandDescriptors,
            SingleRasterSource,
        },
        mock::{MockFeatureCollectionSource, MockRasterSource, MockRasterSourceParams},
        processing::{Expression, ExpressionParams, RasterStacker, RasterStackerParams},
        source::{GdalSource, GdalSourceParameters},
        util::{gdal::add_ndvi_dataset, input::RasterOrVectorOperator},
    };
    use futures::StreamExt;
    use geoengine_datatypes::{
        collections::{ChunksEqualIgnoringCacheHint, MultiPointCollection},
        dataset::NamedData,
        primitives::{
            BandSelection, BoundingBox2D, CacheHint, DateTime, MultiPoint, SpatialPartition2D,
            SpatialResolution, TimeGranularity,
        },
        raster::{
            EmptyGrid2D, GridOrEmpty, RasterDataType, RenameBands, TileInformation,
            TilingSpecification,
        },
        spatial_reference::SpatialReference,
        util::test::TestDefault,
    };

    #[test]
    fn test_ser_de_absolute() {
        let time_shift = TimeShift {
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Raster(
                    GdalSource {
                        params: GdalSourceParameters {
                            data: NamedData::with_system_name("test-raster"),
                        },
                    }
                    .boxed(),
                ),
            },
            params: TimeShiftParams::Absolute {
                time_interval: TimeInterval::new_unchecked(
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                ),
            },
        };

        let serialized = serde_json::to_value(&time_shift).unwrap();

        assert_eq!(
            serialized,
            serde_json::json!({
                "params": {
                    "type": "absolute",
                    "timeInterval": {
                        "start": 1_293_840_000_000_i64,
                        "end": 1_325_376_000_000_i64
                    }
                },
                "sources": {
                    "source": {
                        "type": "GdalSource",
                        "params": {
                            "data": "test-raster"
                        }
                    }
                }
            })
        );

        let deserialized: TimeShift = serde_json::from_value(serialized).unwrap();

        assert_eq!(time_shift.params, deserialized.params);
    }

    #[test]
    fn test_ser_de_relative() {
        let time_shift = TimeShift {
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Raster(
                    GdalSource {
                        params: GdalSourceParameters {
                            data: NamedData::with_system_name("test-raster"),
                        },
                    }
                    .boxed(),
                ),
            },
            params: TimeShiftParams::Relative {
                granularity: TimeGranularity::Years,
                value: 1,
            },
        };

        let serialized = serde_json::to_value(&time_shift).unwrap();

        assert_eq!(
            serialized,
            serde_json::json!({
                "params": {
                    "type": "relative",
                    "granularity": "years",
                    "value": 1
                },
                "sources": {
                    "source": {
                        "type": "GdalSource",
                        "params": {
                            "data": "test-raster"
                        }
                    }
                }
            })
        );

        let deserialized: TimeShift = serde_json::from_value(serialized).unwrap();

        assert_eq!(time_shift.params, deserialized.params);
    }

    #[tokio::test]
    async fn test_absolute_vector_shift() {
        let execution_context = MockExecutionContext::test_default();
        let query_context = MockQueryContext::test_default();

        let source = MockFeatureCollectionSource::single(
            MultiPointCollection::from_data(
                MultiPoint::many(vec![(0., 0.), (1., 1.), (2., 2.)]).unwrap(),
                vec![
                    TimeInterval::new(
                        DateTime::new_utc(2009, 1, 1, 0, 0, 0),
                        DateTime::new_utc_with_millis(2010, 12, 31, 23, 59, 59, 999),
                    )
                    .unwrap(),
                    TimeInterval::new(
                        DateTime::new_utc(2009, 6, 3, 0, 0, 0),
                        DateTime::new_utc(2010, 7, 14, 0, 0, 0),
                    )
                    .unwrap(),
                    TimeInterval::new(
                        DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                        DateTime::new_utc_with_millis(2011, 3, 31, 23, 59, 59, 999),
                    )
                    .unwrap(),
                ],
                Default::default(),
                CacheHint::default(),
            )
            .unwrap(),
        );

        let time_shift = TimeShift {
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Vector(source.boxed()),
            },
            params: TimeShiftParams::Absolute {
                time_interval: TimeInterval::new(
                    DateTime::new_utc(2009, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2009, 6, 1, 0, 0, 0),
                )
                .unwrap(),
            },
        };

        let query_processor = VectorOperator::boxed(time_shift)
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap()
            .query_processor()
            .unwrap()
            .multi_point()
            .unwrap();

        let mut stream = query_processor
            .vector_query(
                VectorQueryRectangle {
                    spatial_bounds: BoundingBox2D::new((0., 0.).into(), (2., 2.).into()).unwrap(),
                    time_interval: TimeInterval::new(
                        DateTime::new_utc(2009, 1, 1, 0, 0, 0),
                        DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                    )
                    .unwrap(),
                    spatial_resolution: SpatialResolution::one(),
                    attributes: ColumnSelection::all(),
                },
                &query_context,
            )
            .await
            .unwrap();

        let mut result = Vec::new();
        while let Some(collection) = stream.next().await {
            result.push(collection.unwrap());
        }

        assert_eq!(result.len(), 1);

        let expected = MultiPointCollection::from_data(
            MultiPoint::many(vec![(0., 0.)]).unwrap(),
            vec![TimeInterval::new(
                DateTime::new_utc(2009, 1, 1, 0, 0, 0),
                DateTime::new_utc_with_millis(2013, 8, 1, 23, 59, 59, 999),
            )
            .unwrap()],
            Default::default(),
            CacheHint::default(),
        )
        .unwrap();

        assert!(result[0].chunks_equal_ignoring_cache_hint(&expected));
    }

    #[tokio::test]
    async fn test_relative_vector_shift() {
        let execution_context = MockExecutionContext::test_default();
        let query_context = MockQueryContext::test_default();

        let source = MockFeatureCollectionSource::single(
            MultiPointCollection::from_data(
                MultiPoint::many(vec![(0., 0.), (1., 1.), (2., 2.)]).unwrap(),
                vec![
                    TimeInterval::new(
                        DateTime::new_utc(2009, 1, 1, 0, 0, 0),
                        DateTime::new_utc_with_millis(2010, 12, 31, 23, 59, 59, 999),
                    )
                    .unwrap(),
                    TimeInterval::new(
                        DateTime::new_utc(2009, 6, 3, 0, 0, 0),
                        DateTime::new_utc(2010, 7, 14, 0, 0, 0),
                    )
                    .unwrap(),
                    TimeInterval::new(
                        DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                        DateTime::new_utc_with_millis(2011, 3, 31, 23, 59, 59, 999),
                    )
                    .unwrap(),
                ],
                Default::default(),
                CacheHint::default(),
            )
            .unwrap(),
        );

        let time_shift = TimeShift {
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Vector(source.boxed()),
            },
            params: TimeShiftParams::Relative {
                granularity: TimeGranularity::Years,
                value: -1,
            },
        };

        let query_processor = VectorOperator::boxed(time_shift)
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap()
            .query_processor()
            .unwrap()
            .multi_point()
            .unwrap();

        let mut stream = query_processor
            .vector_query(
                VectorQueryRectangle {
                    spatial_bounds: BoundingBox2D::new((0., 0.).into(), (2., 2.).into()).unwrap(),
                    time_interval: TimeInterval::new(
                        DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                        DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    )
                    .unwrap(),
                    spatial_resolution: SpatialResolution::one(),
                    attributes: ColumnSelection::all(),
                },
                &query_context,
            )
            .await
            .unwrap();

        let mut result = Vec::new();
        while let Some(collection) = stream.next().await {
            result.push(collection.unwrap());
        }

        assert_eq!(result.len(), 1);

        let expected = MultiPointCollection::from_data(
            MultiPoint::many(vec![(0., 0.), (1., 1.)]).unwrap(),
            vec![
                TimeInterval::new(
                    DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                    DateTime::new_utc_with_millis(2011, 12, 31, 23, 59, 59, 999),
                )
                .unwrap(),
                TimeInterval::new(
                    DateTime::new_utc(2010, 6, 3, 0, 0, 0),
                    DateTime::new_utc(2011, 7, 14, 0, 0, 0),
                )
                .unwrap(),
            ],
            Default::default(),
            CacheHint::default(),
        )
        .unwrap();

        assert!(result[0].chunks_equal_ignoring_cache_hint(&expected));
    }

    #[tokio::test]
    #[allow(clippy::too_many_lines)]
    async fn test_absolute_raster_shift() {
        let empty_grid = GridOrEmpty::Empty(EmptyGrid2D::<u8>::new([3, 2].into()));
        let raster_tiles = vec![
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 0].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 1].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 0].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 1].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2013, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 0].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2013, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 1].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
        ];

        let mrs = MockRasterSource {
            params: MockRasterSourceParams {
                data: raster_tiles,
                result_descriptor: RasterResultDescriptor {
                    data_type: RasterDataType::U8,
                    spatial_reference: SpatialReference::epsg_4326().into(),
                    time: None,
                    bbox: None,
                    resolution: None,
                    bands: RasterBandDescriptors::new_single_band(),
                },
            },
        }
        .boxed();

        let time_shift = TimeShift {
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Raster(mrs),
            },
            params: TimeShiftParams::Absolute {
                time_interval: TimeInterval::new_unchecked(
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                ),
            },
        };

        let execution_context = MockExecutionContext::new_with_tiling_spec(
            TilingSpecification::new((0., 0.).into(), [3, 2].into()),
        );
        let query_context = MockQueryContext::test_default();

        let query_processor = RasterOperator::boxed(time_shift)
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap()
            .query_processor()
            .unwrap()
            .get_u8()
            .unwrap();

        let mut stream = query_processor
            .raster_query(
                RasterQueryRectangle {
                    spatial_bounds: SpatialPartition2D::new_unchecked(
                        (0., 3.).into(),
                        (4., 0.).into(),
                    ),
                    time_interval: TimeInterval::new(
                        DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                        DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    )
                    .unwrap(),
                    spatial_resolution: SpatialResolution::one(),
                    attributes: BandSelection::first(),
                },
                &query_context,
            )
            .await
            .unwrap();

        let mut result = Vec::new();
        while let Some(tile) = stream.next().await {
            result.push(tile.unwrap());
        }

        assert_eq!(result.len(), 2);

        assert_eq!(
            result[0].time,
            TimeInterval::new_unchecked(
                DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                DateTime::new_utc(2011, 1, 1, 0, 0, 0),
            ),
        );
        assert_eq!(
            result[1].time,
            TimeInterval::new_unchecked(
                DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                DateTime::new_utc(2011, 1, 1, 0, 0, 0),
            ),
        );
    }

    #[tokio::test]
    #[allow(clippy::too_many_lines)]
    async fn test_relative_raster_shift() {
        let empty_grid = GridOrEmpty::Empty(EmptyGrid2D::<u8>::new([3, 2].into()));
        let raster_tiles = vec![
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 0].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 1].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 0].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 1].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2013, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 0].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
            RasterTile2D::new_with_tile_info(
                TimeInterval::new_unchecked(
                    DateTime::new_utc(2012, 1, 1, 0, 0, 0),
                    DateTime::new_utc(2013, 1, 1, 0, 0, 0),
                ),
                TileInformation {
                    global_tile_position: [-1, 1].into(),
                    tile_size_in_pixels: [3, 2].into(),
                    global_geo_transform: TestDefault::test_default(),
                },
                0,
                empty_grid.clone(),
                CacheHint::default(),
            ),
        ];

        let mrs = MockRasterSource {
            params: MockRasterSourceParams {
                data: raster_tiles,
                result_descriptor: RasterResultDescriptor {
                    data_type: RasterDataType::U8,
                    spatial_reference: SpatialReference::epsg_4326().into(),
                    time: None,
                    bbox: None,
                    resolution: None,
                    bands: RasterBandDescriptors::new_single_band(),
                },
            },
        }
        .boxed();

        let time_shift = TimeShift {
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Raster(mrs),
            },
            params: TimeShiftParams::Relative {
                granularity: TimeGranularity::Years,
                value: 1,
            },
        };

        let execution_context = MockExecutionContext::new_with_tiling_spec(
            TilingSpecification::new((0., 0.).into(), [3, 2].into()),
        );
        let query_context = MockQueryContext::test_default();

        let query_processor = RasterOperator::boxed(time_shift)
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap()
            .query_processor()
            .unwrap()
            .get_u8()
            .unwrap();

        let mut stream = query_processor
            .raster_query(
                RasterQueryRectangle {
                    spatial_bounds: SpatialPartition2D::new_unchecked(
                        (0., 3.).into(),
                        (4., 0.).into(),
                    ),
                    time_interval: TimeInterval::new(
                        DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                        DateTime::new_utc(2011, 1, 1, 0, 0, 0),
                    )
                    .unwrap(),
                    spatial_resolution: SpatialResolution::one(),
                    attributes: BandSelection::first(),
                },
                &query_context,
            )
            .await
            .unwrap();

        let mut result = Vec::new();
        while let Some(tile) = stream.next().await {
            result.push(tile.unwrap());
        }

        assert_eq!(result.len(), 2);

        assert_eq!(
            result[0].time,
            TimeInterval::new_unchecked(
                DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                DateTime::new_utc(2011, 1, 1, 0, 0, 0),
            ),
        );
        assert_eq!(
            result[1].time,
            TimeInterval::new_unchecked(
                DateTime::new_utc(2010, 1, 1, 0, 0, 0),
                DateTime::new_utc(2011, 1, 1, 0, 0, 0),
            ),
        );
    }

    #[tokio::test]
    async fn test_expression_on_shifted_raster() {
        let mut execution_context = MockExecutionContext::test_default();

        let ndvi_source = GdalSource {
            params: GdalSourceParameters {
                data: add_ndvi_dataset(&mut execution_context),
            },
        }
        .boxed();

        let shifted_ndvi_source = RasterOperator::boxed(TimeShift {
            params: TimeShiftParams::Relative {
                granularity: TimeGranularity::Months,
                value: -1,
            },
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Raster(ndvi_source.clone()),
            },
        });

        let expression = Expression {
            params: ExpressionParams {
                expression: "A - B".to_string(),
                output_type: RasterDataType::F64,
                output_band: None,
                map_no_data: false,
            },
            sources: SingleRasterSource {
                raster: RasterStacker {
                    params: RasterStackerParams {
                        rename_bands: RenameBands::Default,
                    },
                    sources: MultipleRasterSources {
                        rasters: vec![ndvi_source, shifted_ndvi_source],
                    },
                }
                .boxed(),
            },
        }
        .boxed();

        let query_processor = expression
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap()
            .query_processor()
            .unwrap()
            .get_f64()
            .unwrap();

        let query_context = MockQueryContext::test_default();

        let mut stream = query_processor
            .raster_query(
                RasterQueryRectangle {
                    spatial_bounds: SpatialPartition2D::new_unchecked(
                        (-180., 90.).into(),
                        (180., -90.).into(),
                    ),
                    time_interval: TimeInterval::new_instant(DateTime::new_utc(
                        2014, 3, 1, 0, 0, 0,
                    ))
                    .unwrap(),
                    spatial_resolution: SpatialResolution::one(),
                    attributes: BandSelection::first(),
                },
                &query_context,
            )
            .await
            .unwrap();

        let mut result = Vec::new();
        while let Some(tile) = stream.next().await {
            result.push(tile.unwrap());
        }

        assert_eq!(result.len(), 4);
        assert_eq!(
            result[0].time,
            TimeInterval::new(
                DateTime::new_utc(2014, 3, 1, 0, 0, 0),
                DateTime::new_utc(2014, 4, 1, 0, 0, 0)
            )
            .unwrap()
        );
    }

    #[tokio::test]
    async fn test_expression_on_absolute_shifted_raster() {
        let mut execution_context = MockExecutionContext::test_default();

        let ndvi_source = GdalSource {
            params: GdalSourceParameters {
                data: add_ndvi_dataset(&mut execution_context),
            },
        }
        .boxed();

        let shifted_ndvi_source = RasterOperator::boxed(TimeShift {
            params: TimeShiftParams::Absolute {
                time_interval: TimeInterval::new_instant(DateTime::new_utc(2014, 5, 1, 0, 0, 0))
                    .unwrap(),
            },
            sources: SingleRasterOrVectorSource {
                source: RasterOrVectorOperator::Raster(ndvi_source),
            },
        });

        let query_processor = shifted_ndvi_source
            .initialize(WorkflowOperatorPath::initialize_root(), &execution_context)
            .await
            .unwrap()
            .query_processor()
            .unwrap()
            .get_u8()
            .unwrap();

        let query_context = MockQueryContext::test_default();

        let mut stream = query_processor
            .raster_query(
                RasterQueryRectangle {
                    spatial_bounds: SpatialPartition2D::new_unchecked(
                        (-180., 90.).into(),
                        (180., -90.).into(),
                    ),
                    time_interval: TimeInterval::new_instant(DateTime::new_utc(
                        2014, 3, 1, 0, 0, 0,
                    ))
                    .unwrap(),
                    spatial_resolution: SpatialResolution::one(),
                    attributes: BandSelection::first(),
                },
                &query_context,
            )
            .await
            .unwrap();

        let mut result = Vec::new();
        while let Some(tile) = stream.next().await {
            result.push(tile.unwrap());
        }

        assert_eq!(result.len(), 4);
        assert_eq!(
            result[0].time,
            TimeInterval::new(
                DateTime::new_utc(2014, 3, 1, 0, 0, 0),
                DateTime::new_utc(2014, 4, 1, 0, 0, 0)
            )
            .unwrap()
        );
    }

    #[test]
    fn shift_eternal() {
        let eternal = TimeInterval::default();

        let f_shift = RelativeForwardShift {
            step: TimeStep {
                granularity: TimeGranularity::Seconds,
                step: 1,
            },
        };

        let (shifted, state) = f_shift.shift(eternal).unwrap();
        assert_eq!(shifted, eternal);
        let reverse_shifted = f_shift.reverse_shift(eternal, state).unwrap();
        assert_eq!(reverse_shifted, eternal);

        let b_shift = RelativeBackwardShift {
            step: TimeStep {
                granularity: TimeGranularity::Seconds,
                step: 1,
            },
        };

        let (shifted, state) = b_shift.shift(eternal).unwrap();
        assert_eq!(shifted, eternal);

        let reverse_shifted = b_shift.reverse_shift(eternal, state).unwrap();
        assert_eq!(reverse_shifted, eternal);
    }

    #[test]
    fn shift_begin_of_time() {
        let time = TimeInterval::new_unchecked(TimeInstance::MIN, TimeInstance::EPOCH_START);

        let f_shift = RelativeForwardShift {
            step: TimeStep {
                granularity: TimeGranularity::Seconds,
                step: 1,
            },
        };

        let expected_shift =
            TimeInterval::new_unchecked(TimeInstance::MIN, TimeInstance::EPOCH_START + 1_000);

        let (shifted, state) = f_shift.shift(time).unwrap();
        assert_eq!(shifted, expected_shift);

        let reverse_shifted = f_shift.reverse_shift(expected_shift, state).unwrap();
        assert_eq!(reverse_shifted, time);

        let f_shift = RelativeBackwardShift {
            step: TimeStep {
                granularity: TimeGranularity::Seconds,
                step: 1,
            },
        };

        let expected_shift =
            TimeInterval::new_unchecked(TimeInstance::MIN, TimeInstance::EPOCH_START - 1_000);

        let (shifted, state) = f_shift.shift(time).unwrap();
        assert_eq!(shifted, expected_shift);

        let reverse_shifted = f_shift.reverse_shift(expected_shift, state).unwrap();
        assert_eq!(reverse_shifted, time);
    }

    #[test]
    fn shift_end_of_time() {
        let time = TimeInterval::new_unchecked(TimeInstance::EPOCH_START, TimeInstance::MAX);

        let f_shift = RelativeForwardShift {
            step: TimeStep {
                granularity: TimeGranularity::Seconds,
                step: 1,
            },
        };

        let expected_shift =
            TimeInterval::new_unchecked(TimeInstance::EPOCH_START + 1_000, TimeInstance::MAX);

        let (shifted, state) = f_shift.shift(time).unwrap();
        assert_eq!(shifted, expected_shift);

        let reverse_shifted = f_shift.reverse_shift(expected_shift, state).unwrap();
        assert_eq!(reverse_shifted, time);

        let f_shift = RelativeBackwardShift {
            step: TimeStep {
                granularity: TimeGranularity::Seconds,
                step: 1,
            },
        };

        let expected_shift =
            TimeInterval::new_unchecked(TimeInstance::EPOCH_START - 1_000, TimeInstance::MAX);

        let (shifted, state) = f_shift.shift(time).unwrap();
        assert_eq!(shifted, expected_shift);

        let reverse_shifted = f_shift.reverse_shift(expected_shift, state).unwrap();
        assert_eq!(reverse_shifted, time);
    }
}

geo-engine / geoengine / 12469296660

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