diff --git a/openeo_driver/backend.py b/openeo_driver/backend.py index 26a7e445..ec08d46d 100644 --- a/openeo_driver/backend.py +++ b/openeo_driver/backend.py @@ -173,6 +173,12 @@ class LoadParameters(dict): temporal_extent = dict_item(default=(None, None)) spatial_extent = dict_item(default={}) global_extent = dict_item(default={}) # TODO #441 to be removed + # Identifies the source (e.g. `load_collection`/`load_stac` node) these parameters were + # extracted for. Useful (e.g. in post-dry-run global extent determination) to look up + # per-source information that can not simply be merged across all sources + # (e.g. the "mask" branch of a `mask` process potentially needing a different/larger extent + # than the data cube it masks). + source_id = dict_item(default=None) filter_temporal_labels = dict_item(default=None) bands = dict_item(default=None) properties = dict_item(default={}) @@ -213,6 +219,13 @@ def _copy_for_eq(self): copy = super().copy() #load_collection with or without data mask should return an 'equivalent' result copy["data_mask"] = {} + # `source_id` is just an identifier of the process graph node these parameters were + # extracted for. It should not influence equality (e.g. for load_collection caching): + # identical loads from different process graph nodes should be considered equivalent. + # Anything source-specific that influences the actual result (e.g. a per-source + # global extent) should be resolved into other (equality-relevant) fields or the + # EvalEnv before caching boundaries. + copy["source_id"] = None return copy def __eq__(self, o: object) -> bool: diff --git a/openeo_driver/dry_run.py b/openeo_driver/dry_run.py index e894134c..45356963 100644 --- a/openeo_driver/dry_run.py +++ b/openeo_driver/dry_run.py @@ -146,6 +146,7 @@ def _extract_resample_constraint(operation_name: str, args: Union[dict, tuple]) ], ), PropagationRule("aggregate_spatial", operations=["aggregate_spatial"]), + PropagationRule("is_mask", operations=["is_mask"]), PropagationRule("sar_backscatter", operations=["sar_backscatter"]), PropagationRule("process_type", operations=["process_type"]), PropagationRule("custom_cloud_mask", operations=["custom_cloud_mask"]), @@ -658,6 +659,10 @@ def mask(self, mask: "DryRunDataCube", replacement=None) -> "DryRunDataCube": # TODO: if mask cube has no temporal or bbox extent: copy from self? # TODO: or add reference to the self trace to the mask trace and vice versa? mask_resampled = mask._process("resample_cube_spatial", arguments={"target": self, "method": "near"}) + # Tag the mask branch so that downstream (post-dry-run) logic can distinguish + # the "mask" source from the "data" source being masked (e.g. to avoid inflating + # the data cube's extent with the (possibly larger) extent of the mask cube). + mask_resampled = mask_resampled._process("is_mask", arguments={}) cube = self._process("mask", {"mask": mask_resampled}) return DryRunDataCube( traces=cube._traces + mask_resampled._traces, data_tracer=cube._data_tracer, metadata=cube.metadata diff --git a/openeo_driver/processgraph/load_params.py b/openeo_driver/processgraph/load_params.py index 4b3c950f..bd71afd6 100644 --- a/openeo_driver/processgraph/load_params.py +++ b/openeo_driver/processgraph/load_params.py @@ -220,6 +220,7 @@ def _extract_load_parameters(env: EvalEnv, source_id: SourceId) -> LoadParameter source_constraints.remove((source_id, constraints)) # Side effect! params = LoadParameters() + params.source_id = source_id if temporal_extent := constraints.get("temporal_extent"): params.temporal_extent = temporal_extent labels_args = constraints.get("filter_labels", {}) diff --git a/tests/test_dry_run.py b/tests/test_dry_run.py index 99d3d1ee..066167ea 100644 --- a/tests/test_dry_run.py +++ b/tests/test_dry_run.py @@ -319,6 +319,7 @@ def test_evaluate_graph_diamond(dry_run_env, dry_run_tracer): ("load_collection", ("S2_FOOBAR", ()), "load1"), { "bands": ["grass"], + 'is_mask': {}, "resample": {"method": "near", "resolution": (10, 10), "target_crs": CRS_UTM}, "spatial_extent": {"west": 1, "east": 2, "south": 51, "north": 52, "crs": "EPSG:4326"}, }, @@ -2989,7 +2990,7 @@ def test_deduplicate_source_constraints(dry_run_env, dry_run_tracer): assert source_constraints == [ ( ("load_stac", ("https://stac.test/collection1", (), ()), "loadstac1"), - {"temporal_extent": ("2026-01-01", "2026-02-02")}, + {'is_mask': {}, "temporal_extent": ("2026-01-01", "2026-02-02")}, ), ( ("load_stac", ("https://stac.test/collection1", (), ()), "loadstac1"), @@ -3005,8 +3006,12 @@ def test_deduplicate_source_constraints(dry_run_env, dry_run_tracer): assert deduped == [ ( ("load_stac", ("https://stac.test/collection1", (), ()), "loadstac1"), - {"temporal_extent": ("2026-01-01", "2026-02-02")}, + {'is_mask': {},"temporal_extent": ("2026-01-01", "2026-02-02")}, ), + ( + ("load_stac", ("https://stac.test/collection1", (), ()), "loadstac1"), + {"temporal_extent": ("2026-01-01", "2026-02-02")}, + ) ]