feat(load): add arrow endpoints

antarest/core/utils/string.py

@@ -18,3 +18,7 @@ def to_pascal_case(value: str) -> str:
 def to_camel_case(value: str) -> str:
     v = to_pascal_case(value)
     return v[0].lower() + v[1:] if len(v) > 0 else ""
+
+
+def to_kebab_case(string: str) -> str:
+    return string.replace("_", "-")

antarest/study/business/load_management.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2024, RTE (https://www.rte-france.com)
+#
+# See AUTHORS.txt
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#
+# SPDX-License-Identifier: MPL-2.0
+#
+# This file is part of the Antares project.
+import io
+from typing import cast
+
+import pandas as pd
+from starlette.responses import JSONResponse, Response
+
+from antarest.core.model import JSON
+from antarest.study.business.model.load_model import LoadDTO
+from antarest.study.model import MatrixFormat, Study
+from antarest.study.storage.rawstudy.model.filesystem.matrix.input_series_matrix import InputSeriesMatrix
+from antarest.study.storage.storage_service import StudyStorageService
+
+LOAD_PATH = "input/load/series/load_{area_id}"
+
+
+class LoadManager:
+    def __init__(self, storage_service: StudyStorageService) -> None:
+        self.storage_service = storage_service
+
+    def get_load_matrix(self, study: Study, area_id: str, matrix_format: MatrixFormat) -> Response:
+        file_study = self.storage_service.get_storage(study).get_raw(study)
+
+        load_path = LOAD_PATH.format(area_id=area_id).split("/")
+        node = file_study.tree.get_node(load_path)
+
+        if not isinstance(node, InputSeriesMatrix):
+            return Response(content="Invalid node type", status_code=400)
+
+        matrix_data = InputSeriesMatrix.parse(node, return_dataframe=(matrix_format == MatrixFormat.ARROW))
+
+        if matrix_format == MatrixFormat.JSON:
+            return JSONResponse(content=cast(JSON, matrix_data))
+
+        if matrix_format == MatrixFormat.ARROW:
+            matrix_df = cast(pd.DataFrame, matrix_data)
+            matrix_df.columns = matrix_df.columns.map(str)
+            buffer = io.BytesIO()
+            matrix_df.to_feather(buffer, compression="uncompressed")
+            return Response(content=buffer.getvalue(), media_type="application/octet-stream")
+
+        return Response(content="Unsupported format", status_code=422)
+
+    def update_load_matrix(self, study: Study, area_id: str, load_dto: LoadDTO) -> LoadDTO:
+        load_properties = load_dto.to_properties()
+        load_path = LOAD_PATH.format(area_id=area_id).split("/")
+
+        file_study = self.storage_service.get_storage(study).get_raw(study)
+
+        file_study.tree.save(load_properties.matrix, load_path)
+
+        return load_dto

antarest/study/business/model/load_model.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2024, RTE (https://www.rte-france.com)
+#
+# See AUTHORS.txt
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#
+# SPDX-License-Identifier: MPL-2.0
+#
+# This file is part of the Antares project.
+from pydantic import ConfigDict
+
+from antarest.core.model import JSON
+from antarest.core.serialization import AntaresBaseModel
+from antarest.core.utils.string import to_camel_case, to_kebab_case
+
+
+class LoadDTO(AntaresBaseModel):
+    matrix: JSON | bytes
+
+    model_config = ConfigDict(alias_generator=to_camel_case, extra="forbid")
+
+    def to_properties(self) -> "LoadProperties":
+        return LoadProperties(matrix=self.matrix)
+
+
+class LoadProperties(AntaresBaseModel):
+    matrix: JSON | bytes
+
+    model_config = ConfigDict(alias_generator=to_kebab_case, extra="forbid")
+
+    def to_dto(self) -> LoadDTO:
+        return LoadDTO(matrix=self.matrix)

antarest/study/model.py

@@ -16,6 +16,7 @@
 import typing as t
 import uuid
 from datetime import datetime, timedelta
+from enum import StrEnum
 from pathlib import Path
 
 from antares.study.version import StudyVersion
@@ -525,6 +526,11 @@ def suffix(self) -> str:
         return mapping[self]
 
 
+class MatrixFormat(StrEnum):
+    JSON = "json"
+    ARROW = "arrow"
+
+
 class StudyDownloadDTO(AntaresBaseModel):
     """
     DTO used to download outputs

antarest/study/service.py

@@ -86,6 +86,7 @@
 from antarest.study.business.district_manager import DistrictManager
 from antarest.study.business.general_management import GeneralManager
 from antarest.study.business.link_management import LinkInfoDTO, LinkManager
+from antarest.study.business.load_management import LoadManager
 from antarest.study.business.matrix_management import MatrixManager, MatrixManagerError
 from antarest.study.business.optimization_management import OptimizationManager
 from antarest.study.business.playlist_management import PlaylistManager
@@ -365,6 +366,7 @@ def __init__(
         self.adequacy_patch_manager = AdequacyPatchManager(self.storage_service)
         self.advanced_parameters_manager = AdvancedParamsManager(self.storage_service)
         self.hydro_manager = HydroManager(self.storage_service)
+        self.load_manager = LoadManager(self.storage_service)
         self.allocation_manager = AllocationManager(self.storage_service)
         self.properties_manager = PropertiesManager(self.storage_service)
         self.renewable_manager = RenewableManager(self.storage_service)

antarest/study/web/study_data_blueprint.py

@@ -17,7 +17,7 @@
 
 import typing_extensions as te
 from fastapi import APIRouter, Body, Depends, Query
-from starlette.responses import RedirectResponse
+from starlette.responses import RedirectResponse, Response
 
 from antarest.core.config import Config
 from antarest.core.jwt import JWTUser
@@ -69,13 +69,14 @@
 from antarest.study.business.district_manager import DistrictCreationDTO, DistrictInfoDTO, DistrictUpdateDTO
 from antarest.study.business.general_management import GeneralFormFields
 from antarest.study.business.link_management import LinkInfoDTO
+from antarest.study.business.model.load_model import LoadDTO
 from antarest.study.business.optimization_management import OptimizationFormFields
 from antarest.study.business.playlist_management import PlaylistColumns
 from antarest.study.business.scenario_builder_management import Rulesets, ScenarioType
 from antarest.study.business.table_mode_management import TableDataDTO, TableModeType
 from antarest.study.business.thematic_trimming_field_infos import ThematicTrimmingFormFields
 from antarest.study.business.timeseries_config_management import TSFormFields
-from antarest.study.model import PatchArea, PatchCluster
+from antarest.study.model import MatrixFormat, PatchArea, PatchCluster
 from antarest.study.service import StudyService
 from antarest.study.storage.rawstudy.model.filesystem.config.binding_constraint import (
     BindingConstraintFrequency,
@@ -524,6 +525,40 @@ def update_inflow_structure(
         study = study_service.check_study_access(uuid, StudyPermissionType.WRITE, params)
         study_service.hydro_manager.update_inflow_structure(study, area_id, values)
 
+    @bp.get(
+        "/{uuid}/{area_id}/load/series",
+        tags=[APITag.study_data],
+        summary="Get load series data",
+    )
+    def get_load_series(
+        uuid: str,
+        area_id: str,
+        matrix_format: MatrixFormat = MatrixFormat.JSON,
+        current_user: JWTUser = Depends(auth.get_current_user),
+    ) -> Response:
+        logger.info(
+            msg=f"Getting load series data for area {area_id} of study {uuid}",
+            extra={"user": current_user.id},
+        )
+        params = RequestParameters(user=current_user)
+        study = study_service.check_study_access(uuid, StudyPermissionType.READ, params)
+        return study_service.load_manager.get_load_matrix(study, area_id, matrix_format)
+
+    @bp.put(
+        "/{uuid}/{area_id}/load/series",
+        tags=[APITag.study_data],
+        summary="Update load series data",
+    )
+    def update_load_series(
+        uuid: str,
+        area_id: str,
+        load_dto: LoadDTO,
+        current_user: JWTUser = Depends(auth.get_current_user),
+    ) -> LoadDTO:
+        params = RequestParameters(user=current_user)
+        study = study_service.check_study_access(uuid, StudyPermissionType.WRITE, params)
+        return study_service.load_manager.update_load_matrix(study, area_id, load_dto)
+
     @bp.put(
         "/studies/{uuid}/matrix",
         tags=[APITag.study_data],

requirements.txt

@@ -38,6 +38,7 @@ pandas~=2.2.3
 paramiko~=3.4.1
 plyer~=2.0.0
 psycopg2-binary~=2.9.9
+pyarrow~=18.0.0
 py7zr~=0.20.6
 python-json-logger~=2.0.7
 PyYAML~=5.3.1

tests/integration/study_data_blueprint/test_load.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2024, RTE (https://www.rte-france.com)
+#
+# See AUTHORS.txt
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#
+# SPDX-License-Identifier: MPL-2.0
+#
+# This file is part of the Antares project.
+import pandas as pd
+import pytest
+from starlette.testclient import TestClient
+
+from tests.integration.prepare_proxy import PreparerProxy
+
+
+@pytest.mark.unit_test
+class TestLink:
+    @pytest.mark.parametrize("study_type", ["raw", "variant"])
+    def test_load(self, client: TestClient, user_access_token: str, study_type: str) -> None:
+        client.headers = {"Authorization": f"Bearer {user_access_token}"}  # type: ignore
+
+        preparer = PreparerProxy(client, user_access_token)
+        study_id = preparer.create_study("foo", version=880)
+
+        if study_type == "variant":
+            study_id = preparer.create_variant(study_id, name="Variant 1")
+
+        area1_id = preparer.create_area(study_id, name="Area 1")["id"]
+
+        # Test simple get JSON
+
+        res = client.get(f"/v1/{study_id}/{area1_id}/load/series?matrix_format=json")
+        assert res.status_code == 200
+        assert res.headers["content-type"] == "application/json"
+        assert pd.DataFrame(data=res.json()["data"], columns=res.json()["columns"]).shape == (8760, 1)
+
+        # Test simple get ARROW
+
+        res = client.get(f"/v1/{study_id}/{area1_id}/load/series?matrix_format=arrow")
+        assert res.status_code == 200
+        assert res.headers["content-type"] == "application/octet-stream"
+
+        # Test get wrong format
+
+        res = client.get(f"/v1/{study_id}/{area1_id}/load/series?matrix_format=wrong_format")
+        assert res.status_code == 422
+        expected = {
+            "body": None,
+            "description": "Input should be 'json' or 'arrow'",
+            "exception": "RequestValidationError",
+        }
+        assert res.json() == expected