saveRDDAllowAssetPerBand: handle items instead of files #1111

Author: bossie

openeogeotrellis/geopysparkdatacube.py

@@ -2108,7 +2108,7 @@ def add_gdalinfo_objects(assets_original):
                                 }
                             )
                         else:
-                            paths_tuples = get_jvm().org.openeo.geotrellis.geotiff.package.saveRDDAllowAssetPerBand(
+                            java_items = get_jvm().org.openeo.geotrellis.geotiff.package.saveRDDAllowAssetPerBand(
                                 max_level_rdd,
                                 band_count,
                                 str(save_filename),
@@ -2117,21 +2117,40 @@ def add_gdalinfo_objects(assets_original):
                                 gtiff_options,
                             )
 
-                            assets = {}
-                            for path, bbox, band_indices in map(to_tuple, paths_tuples):
-                                file_name = str(pathlib.Path(path).relative_to(save_directory))
-                                assets[file_name] = {
-                                    "href": str(path),
-                                    "type": "image/tiff; application=geotiff",
-                                    "roles": ["data"],
-                                    "bands": [band for i, band in enumerate(bands) if i in band_indices],
-                                    "nodata": nodata,
+                            items = {}
+                            for java_item in java_items:
+                                assets = {}
+
+                                bbox = java_item.bbox()
+
+                                for asset_key, asset in java_item.assets().items():
+                                    path = asset.path()
+                                    band_indices = asset.bandIndices()
+
+                                    assets[asset_key] = {
+                                        "href": str(path),
+                                        "type": "image/tiff; application=geotiff",
+                                        "roles": ["data"],
+                                        "bands": [band for i, band in enumerate(bands) if i in band_indices],
+                                        "nodata": nodata,
+                                    }
+                                    if bbox:
+                                        assets[asset_key]["bbox"] = to_latlng_bbox(bbox)
+                                        assets[asset_key]["geometry"] = mapping(
+                                            Polygon.from_bounds(*to_latlng_bbox(bbox))
+                                        )
+
+                                assets = add_gdalinfo_objects(assets)
+                                item = {
+                                    "id": java_item.id(),
+                                    "geometry": mapping(Polygon.from_bounds(*to_latlng_bbox(bbox))),
+                                    "bbox": to_latlng_bbox(bbox),
+                                    "assets": assets,
                                 }
-                                if bbox:
-                                    assets[file_name]["bbox"] = to_latlng_bbox(bbox)
-                                    assets[file_name]["geometry"] = mapping(Polygon.from_bounds(*to_latlng_bbox(bbox)))
 
-                            return add_gdalinfo_objects(assets)
+                                items[java_item.id()] = item
+
+                            return items
 
             else:
                 if not save_filename.endswith(".png"):

tests/test_batch_result.py

@@ -2624,8 +2624,10 @@ def apply_datacube(cube: XarrayDataCube, context: dict) -> XarrayDataCube:
         ),
     ],
 )
-def test_unified_asset_keys(tmp_path, separate_asset_per_band, expected_tiff_files, expected_asset_keys):
-    process_graph = {  # plain old spatiotemporal data cube to GeoTIFF
+def test_unified_asset_keys_spatiotemporal_geotiff(
+    tmp_path, separate_asset_per_band, expected_tiff_files, expected_asset_keys
+):
+    process_graph = {
         "load2": {
             "process_id": "load_collection",
             "arguments": {
@@ -2842,3 +2844,100 @@ def test_unified_asset_keys_sample_by_feature(tmp_path):
 
     for item in items:
         assert set(item["assets"].keys()) == {"openEO"}
+
+
+@pytest.mark.parametrize(
+    ["separate_asset_per_band", "expected_tiff_files", "expected_asset_keys"],
+    [
+        (False, {"openEO.tif"}, {"openEO"}),
+        (
+            True,
+            {
+                "openEO_Flat:0.tif",
+                "openEO_Flat:1.tif",
+                "openEO_Flat:2.tif",
+            },
+            {"openEO_Flat:0", "openEO_Flat:1", "openEO_Flat:2"},
+        ),
+    ],
+)
+def test_unified_asset_keys_spatial_geotiff(
+    tmp_path, separate_asset_per_band, expected_tiff_files, expected_asset_keys
+):
+    process_graph = {
+        "load2": {
+            "process_id": "load_collection",
+            "arguments": {
+                "bands": [
+                    "Flat:0",
+                    "Flat:1",
+                    "Flat:2",
+                ],
+                "id": "TestCollection-LonLat16x16",
+                "spatial_extent": {
+                    "west": 0,
+                    "south": 50,
+                    "east": 5,
+                    "north": 55,
+                },
+                "temporal_extent": ["2025-04-01", "2025-04-21"],
+            },
+        },
+        "reducedimension1": {
+            "process_id": "reduce_dimension",
+            "arguments": {
+                "data": {"from_node": "load2"},
+                "dimension": "t",
+                "reducer": {
+                    "process_graph": {
+                        "first1": {
+                            "process_id": "first",
+                            "arguments": {"data": {"from_parameter": "data"}},
+                            "result": True,
+                        }
+                    }
+                },
+            },
+        },
+        "save1": {
+            "process_id": "save_result",
+            "arguments": {
+                "data": {"from_node": "reducedimension1"},
+                "format": "GTIFF",
+                "options": {"separate_asset_per_band": separate_asset_per_band},
+            },
+            "result": True,
+        },
+    }
+
+    process = {
+        "process_graph": process_graph,
+    }
+
+    job_dir = tmp_path
+    metadata_file = job_dir / "job_metadata.json"
+
+    run_job(
+        process,
+        output_file=job_dir / "out",
+        metadata_file=metadata_file,
+        api_version="2.0.0",
+        job_dir=job_dir,
+        dependencies=[],
+    )
+
+    tiff_files = {file for file in os.listdir(job_dir) if file.endswith(".tif")}
+    assert tiff_files == expected_tiff_files
+
+    with open(metadata_file) as f:
+        job_metadata = json.load(f)
+
+    items = job_metadata["items"]
+    print(f"items={json.dumps(items, indent=2)}")
+
+    assert len(items) == 1
+    # single item ID can be anything (no spatial or temporal references)
+    assert job_metadata["start_datetime"] == "2025-04-01T00:00:00Z"  # top-level rather than on Item
+    assert job_metadata["end_datetime"] == "2025-04-21T00:00:00Z"  # ditto
+
+    assert set(items[0]["assets"].keys()) == expected_asset_keys