From 7a9315e01136153810cb0ea10e2be2e63b964725 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Thu, 20 Aug 2026 17:57:03 +0200 Subject: [PATCH 01/29] first pieces of implementation: - DB models for variables and areas - functions for extracting them from file outputs to be continued with: - some tests - implement the download API for just areas for now - implement other element types Signed-off-by: Sylvain Leclerc --- antarest/output/filestudy/model.py | 102 +++++++++++- antarest/output/storage/v2/dbmodel.py | 118 ++++++++++++++ .../output/storage/v2/variables_parsing.py | 150 ++++++++++++++++++ 3 files changed, 369 insertions(+), 1 deletion(-) create mode 100644 antarest/output/storage/v2/dbmodel.py create mode 100644 antarest/output/storage/v2/variables_parsing.py diff --git a/antarest/output/filestudy/model.py b/antarest/output/filestudy/model.py index d30f64d1a6..e53090a910 100644 --- a/antarest/output/filestudy/model.py +++ b/antarest/output/filestudy/model.py @@ -11,12 +11,14 @@ # This file is part of the Antares project. from dataclasses import dataclass from enum import Enum, StrEnum +from functools import cached_property from pathlib import Path -from typing import Callable, Generic, Literal, Sequence, TypeAlias, TypeVar +from typing import Callable, Generic, Iterable, Literal, Sequence, TypeAlias, TypeVar import polars as pl from antarest.core.exceptions import OutputSubFolderNotFound +from antarest.study.model import MatrixFrequency """Column name for the Monte Carlo year.""" MCYEAR_COL = "mcYear" @@ -163,3 +165,101 @@ def find_mode_dir(output_dir: Path) -> Path: if mode_dir.exists(): return mode_dir raise OutputSubFolderNotFound(output_dir.name, "economy|adequacy") + + +class FileOutput: + """ + Provides a collection of methods to inspect the content of an output directory. + + Caches properties that may take some time to build (typically going through files), for efficiency. + + Attributes: + output_dir (Path): The path to the output directory typically (/outputs/). + """ + + def __init__(self, output_dir: Path): + self.output_dir = output_dir + + @cached_property + def mc_years(self) -> list[int]: + mode_dir = find_mode_dir(self.output_dir) + mc_ind_dir = mode_dir / "mc-ind" + if not mc_ind_dir.exists(): + return [] + return sorted(int(d.name) for d in mc_ind_dir.iterdir()) + + @property + def first_mc_year(self) -> int: + return self.mc_years[0] + + @property + def mode(self) -> str: + return self.mode_dir.name + + @cached_property + def mode_dir(self) -> Path: + return find_mode_dir(self.output_dir) + + @property + def mc_all_dir(self) -> Path: + return self.mode_dir / "mc-all" + + @property + def mc_ind_dir(self) -> Path: + return self.mode_dir / "mc-ind" + + def get_mc_year_dir(self, year: int) -> Path: + return self.mc_ind_dir / f"{year:05d}" + + @cached_property + def link_ids(self) -> tuple[str, ...]: + """ + IDs of links that have data in the output directory, sorted. + """ + return tuple(sorted(d.name for d in self.iter_links_dir(self.first_mc_year))) + + @cached_property + def area_ids(self) -> tuple[str, ...]: + """ + IDs of areas that have data in the output directory, sorted. + """ + return tuple(sorted(d.name for d in self.iter_areas_dir(self.first_mc_year))) + + def iter_areas_dir(self, mc_year: int) -> Iterable[Path]: + """ + No ordering guarantee. + """ + return (self.get_mc_year_dir(mc_year) / "areas").iterdir() + + def iter_links_dir(self, mc_year: int) -> Iterable[Path]: + """ + No ordering guarantee. + """ + return (self.get_mc_year_dir(mc_year) / "links").iterdir() + + def get_mc_all_file( + self, + file_type: MCAllAreasQueryFile | MCAllLinksQueryFile, + area_id: str, + frequency: MatrixFrequency, + ) -> Path | None: + """ + Returns the path corresponding to the specified data, if it exists. + """ + element_type = "areas" if isinstance(file_type, MCIndAreasQueryFile) else "links" + file_path = self.mc_all_dir / element_type / area_id / f"{file_type}-{frequency}.txt" + return file_path if file_path.exists() else None + + def get_mc_ind_file( + self, + mc_year: int, + file_type: MCIndAreasQueryFile | MCIndLinksQueryFile, + area_id: str, + frequency: MatrixFrequency, + ) -> Path | None: + """ + Returns the path corresponding to the specified data, if it exists. + """ + element_type = "areas" if isinstance(file_type, MCIndAreasQueryFile) else "links" + file_path = self.get_mc_year_dir(mc_year) / element_type / area_id / f"{file_type}-{frequency}.txt" + return file_path if file_path.exists() else None diff --git a/antarest/output/storage/v2/dbmodel.py b/antarest/output/storage/v2/dbmodel.py new file mode 100644 index 0000000000..fd8ce3ee36 --- /dev/null +++ b/antarest/output/storage/v2/dbmodel.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026, 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 typing import Any, Iterable, Literal, TypeAlias + +from sqlalchemy import BigInteger, Dialect, ForeignKeyConstraint, SmallInteger, String, select, types +from sqlalchemy.orm import Mapped, Session, mapped_column +from typing_extensions import override + +from antarest.dbmodel import Base +from antarest.output.filestudy.model import VariableDescription + +ElementType: TypeAlias = Literal[ + "area", + "link", + "binding_constraint", + "thermal_cluster", + "renewable_cluster", + "short_term_storage", +] + +ScenarioAggregation: TypeAlias = Literal["mc-ind", "mc-all"] + + +class DbParquetOutput(Base): + # TODO: we should merge the existing v2_output_metadata tables into this one + + __tablename__ = "parquet_output" + + id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True) + + +class DbParquetVariable(Base): + __tablename__ = "parquet_variable" + + __table_args__ = (ForeignKeyConstraint(["output_id"], ["parquet_output.id"]),) # TODO + + output_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) + scenario_aggregation: Mapped[ScenarioAggregation] = mapped_column(primary_key=True) + element_type: Mapped[ElementType] = mapped_column(primary_key=True) + column: Mapped[int] = mapped_column(SmallInteger, primary_key=True) + name: Mapped[str] + unit: Mapped[str | None] + statistic_type: Mapped[str | None] + + +class Columns(types.TypeDecorator[list[int]]): + """ + Stores a list of columns as a comma separated string. + + Avoids a many to many relationship which would not be useful. + """ + + impl = String + cache_ok = True + + @override + def process_bind_param(self, value: list[int] | None, dialect: Dialect) -> str: + if not isinstance(value, list): + raise ValueError("Expected a list of int for variable columns") + return ",".join(str(c) for c in value) + + @override + def process_result_value(self, value: Any | None, dialect: Dialect) -> list[int]: + if not isinstance(value, str): + raise ValueError("Expected a string in variable columns.") + return [int(c) for c in value.split(",")] + + +class DbParquetArea(Base): + __tablename__ = "parquet_area" + + __table_args__ = ForeignKeyConstraint(["output_id"], ["parquet_output.id"]) # TODO + + output_id: Mapped[int] = mapped_column(BigInteger) + area_id: Mapped[str] + mc_all_vars: Mapped[list[int]] = mapped_column(Columns) + mc_ind_vars: Mapped[list[int]] = mapped_column(Columns) + + +class VariablesIndex: + def __init__(self, variables: Iterable[DbParquetVariable]) -> None: + self._variables: dict[tuple[ScenarioAggregation, ElementType], list[VariableDescription]] = {} + for v in variables: + self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(to_var_model(v)) + + def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: + return self._variables.get((aggregation, element_type), []) + + +def to_var_model(db_var: DbParquetVariable) -> VariableDescription: + return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) + + +def get_area_variables( + session: Session, output_id: int, aggregation: ScenarioAggregation, area_id: str +) -> list[VariableDescription]: + + # All variables, should load fast ? + output_variables = session.execute( + select(DbParquetVariable).where(DbParquetVariable.output_id == output_id) + ).scalars() + variables_index = VariablesIndex(output_variables) + + # Get area information + area = session.execute(select(DbParquetArea).where(DbParquetArea.area_id == area_id)).scalar_one() + + all_areas_vars = variables_index.get_variables(aggregation, "area") + cols = area.mc_all_vars if aggregation == "mc-all" else area.mc_ind_vars + return [all_areas_vars[c] for c in cols] diff --git a/antarest/output/storage/v2/variables_parsing.py b/antarest/output/storage/v2/variables_parsing.py new file mode 100644 index 0000000000..3d5d012217 --- /dev/null +++ b/antarest/output/storage/v2/variables_parsing.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026, 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. + +""" +Extraction of variables metadata from file studies, in order to populate the database +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Callable + +from sqlalchemy.orm import Session + +from antarest.output.filestudy.matrixfiles import get_start_column, parse_headers +from antarest.output.filestudy.model import FileOutput, MCAllAreasQueryFile, MCIndAreasQueryFile, VariableDescription +from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ScenarioAggregation +from antarest.study.model import MatrixFrequency + +# TODO: Possibly a better naming to find than "parsing results" + + +@dataclass(frozen=True) +class ParsingResultPart: + """ + A list of variables, and for each area the list of variables, as a list of indices. + + Results from the parsing of a set of files for one element type, and either mc-ind or mc-all + """ + + variables: list[VariableDescription] + area_vars: dict[str, list[int]] + + +@dataclass(frozen=True) +class OutputParsingResult: + """ + Intermediate data structure from which we'll populate the database. + """ + + mc_ind_areas: ParsingResultPart + mc_all_areas: ParsingResultPart + + +def parse_area_variables(file_output: FileOutput, aggregation: ScenarioAggregation) -> ParsingResultPart: + var_cols: dict[VariableDescription, int] = {} + vars: list[VariableDescription] = [] + area_cols: dict[str, list[int]] = {} + + get_file: Callable[[str, MatrixFrequency], Path | None] + match aggregation: + case "mc-ind": + + def get_file(element_id: str, freq: MatrixFrequency) -> Path | None: + return file_output.get_mc_ind_file( + file_output.first_mc_year, MCIndAreasQueryFile.VALUES, element_id, freq + ) + case "mc-all": + + def get_file(element_id: str, freq: MatrixFrequency) -> Path | None: + return file_output.get_mc_all_file(MCAllAreasQueryFile.VALUES, element_id, freq) + + for element_id in file_output.area_ids: + # searching for the first existing "frequency" + for freq in MatrixFrequency: + if data_file := get_file(element_id, freq): + with open(data_file) as f: + area_vars = parse_headers(f, get_start_column(freq)) + + for v in area_vars: + if v not in var_cols: + var_cols[v] = len(vars) + vars.append(v) + + area_cols[element_id] = [var_cols[v] for v in area_vars] + break # other frequencies will have the same variables + + return ParsingResultPart(variables=vars, area_vars=area_cols) + + +def parse_output_variables(file_output: FileOutput) -> OutputParsingResult: + """ + Extract area "values" variables from the output + """ + + return OutputParsingResult( + mc_all_areas=parse_area_variables(file_output, "mc-all"), + mc_ind_areas=parse_area_variables(file_output, "mc-ind"), + ) + + +def extract_output_variables_to_database(session: Session, output_id: int, file_output: FileOutput) -> None: + """ + Parses variables from file output an dump them to database. + """ + parsing_result = parse_output_variables(file_output) + + variables: list[DbParquetVariable] = [] + areas: list[DbParquetArea] = [] + + mc_all_areas = parsing_result.mc_all_areas + mc_ind_areas = parsing_result.mc_ind_areas + for c, v in enumerate(mc_all_areas.variables): + variables.append( + DbParquetVariable( + output_id=output_id, + scenario_aggregation="mc-all", + element_type="area", + column=c, + name=v.name, + unit=v.unit, + statistic_type=v.statistic_type, + ) + ) + for c, v in enumerate(mc_ind_areas.variables): + variables.append( + DbParquetVariable( + output_id=output_id, + scenario_aggregation="mc-ind", + element_type="area", + column=c, + name=v.name, + unit=v.unit, + statistic_type=v.statistic_type, + ) + ) + + mc_ind_area_vars = mc_ind_areas.area_vars + mc_all_area_vars = mc_all_areas.area_vars + area_ids = sorted(set(mc_all_area_vars).union(mc_ind_area_vars)) + for area_id in area_ids: + areas.append( + DbParquetArea( + output_id=output_id, + area_id=area_id, + mc_all_vars=mc_all_area_vars.get(area_id, []), + mc_ind_vars=mc_ind_area_vars.get(area_id, []), + ) + ) + + session.add_all(variables) + session.add_all(areas) From 15d8257454c179f890175b3985082ba95af948d0 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 12:11:18 +0200 Subject: [PATCH 02/29] unit test for variables parsing Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/dbmodel.py | 6 +-- .../output/storage/v2/variables_parsing.py | 45 ++++++++-------- tests/conftest.py | 5 ++ tests/output/conftest.py | 19 +++++++ .../about-the-study/areas.txt | 0 .../about-the-study/comments.txt | 0 .../about-the-study/links.txt | 0 .../about-the-study/parameters.ini | 0 .../about-the-study/study.ini | 0 .../annualSystemCost.txt | 0 .../checkIntegrity.txt | 0 .../areas/@ all areas/values-monthly.txt | 0 .../mc-ind/00001/areas/es/values-monthly.txt | 0 .../mc-ind/00001/areas/fr/values-monthly.txt | 0 .../areas/@ all areas/values-monthly.txt | 0 .../mc-ind/00002/areas/es/values-monthly.txt | 0 .../mc-ind/00002/areas/fr/values-monthly.txt | 0 .../execution_info.ini | 0 .../info.antares-output | 0 ...t.py => test_matrix_aggregation_result.py} | 11 ++-- .../storage/v2/test_variables_parsing.py | 51 +++++++++++++++++++ 21 files changed, 102 insertions(+), 35 deletions(-) create mode 100644 tests/output/conftest.py rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/about-the-study/areas.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/about-the-study/comments.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/about-the-study/links.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/about-the-study/parameters.ini (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/about-the-study/study.ini (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/annualSystemCost.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/checkIntegrity.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/@ all areas/values-monthly.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/es/values-monthly.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/fr/values-monthly.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/@ all areas/values-monthly.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/es/values-monthly.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/fr/values-monthly.txt (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/execution_info.ini (100%) rename tests/output/{filestudy => }/data/20260810-1420eco-thermal_groups/info.antares-output (100%) rename tests/output/filestudy/{test_matrix_aggregattion_result.py => test_matrix_aggregation_result.py} (87%) create mode 100644 tests/output/storage/v2/test_variables_parsing.py diff --git a/antarest/output/storage/v2/dbmodel.py b/antarest/output/storage/v2/dbmodel.py index fd8ce3ee36..0d0f0936e3 100644 --- a/antarest/output/storage/v2/dbmodel.py +++ b/antarest/output/storage/v2/dbmodel.py @@ -78,10 +78,10 @@ def process_result_value(self, value: Any | None, dialect: Dialect) -> list[int] class DbParquetArea(Base): __tablename__ = "parquet_area" - __table_args__ = ForeignKeyConstraint(["output_id"], ["parquet_output.id"]) # TODO + __table_args__ = (ForeignKeyConstraint(["output_id"], ["parquet_output.id"]),) # TODO - output_id: Mapped[int] = mapped_column(BigInteger) - area_id: Mapped[str] + output_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) + area_id: Mapped[str] = mapped_column(primary_key=True) mc_all_vars: Mapped[list[int]] = mapped_column(Columns) mc_ind_vars: Mapped[list[int]] = mapped_column(Columns) diff --git a/antarest/output/storage/v2/variables_parsing.py b/antarest/output/storage/v2/variables_parsing.py index 3d5d012217..874da9558d 100644 --- a/antarest/output/storage/v2/variables_parsing.py +++ b/antarest/output/storage/v2/variables_parsing.py @@ -22,7 +22,7 @@ from antarest.output.filestudy.matrixfiles import get_start_column, parse_headers from antarest.output.filestudy.model import FileOutput, MCAllAreasQueryFile, MCIndAreasQueryFile, VariableDescription -from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ScenarioAggregation +from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ElementType, ScenarioAggregation from antarest.study.model import MatrixFrequency # TODO: Possibly a better naming to find than "parsing results" @@ -97,6 +97,23 @@ def parse_output_variables(file_output: FileOutput) -> OutputParsingResult: ) +def _convert_to_db_vars( + output_id: int, aggregation: ScenarioAggregation, elt_type: ElementType, vars: list[VariableDescription] +) -> list[DbParquetVariable]: + return [ + DbParquetVariable( + output_id=output_id, + scenario_aggregation=aggregation, + element_type=elt_type, + column=c, + name=v.name, + unit=v.unit, + statistic_type=v.statistic_type, + ) + for c, v in enumerate(vars) + ] + + def extract_output_variables_to_database(session: Session, output_id: int, file_output: FileOutput) -> None: """ Parses variables from file output an dump them to database. @@ -108,30 +125,8 @@ def extract_output_variables_to_database(session: Session, output_id: int, file_ mc_all_areas = parsing_result.mc_all_areas mc_ind_areas = parsing_result.mc_ind_areas - for c, v in enumerate(mc_all_areas.variables): - variables.append( - DbParquetVariable( - output_id=output_id, - scenario_aggregation="mc-all", - element_type="area", - column=c, - name=v.name, - unit=v.unit, - statistic_type=v.statistic_type, - ) - ) - for c, v in enumerate(mc_ind_areas.variables): - variables.append( - DbParquetVariable( - output_id=output_id, - scenario_aggregation="mc-ind", - element_type="area", - column=c, - name=v.name, - unit=v.unit, - statistic_type=v.statistic_type, - ) - ) + variables.extend(_convert_to_db_vars(output_id, "mc-all", "area", mc_all_areas.variables)) + variables.extend(_convert_to_db_vars(output_id, "mc-ind", "area", mc_ind_areas.variables)) mc_ind_area_vars = mc_ind_areas.area_vars mc_all_area_vars = mc_all_areas.area_vars diff --git a/tests/conftest.py b/tests/conftest.py index 5e5d0e1a3e..2d00c9b426 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -68,6 +68,11 @@ def project_path() -> Path: return PROJECT_DIR +@pytest.fixture(scope="session") +def test_root() -> Path: + return HERE + + @pytest.fixture def ini_cleaner() -> Callable[[str], str]: def cleaner(txt: str) -> str: diff --git a/tests/output/conftest.py b/tests/output/conftest.py new file mode 100644 index 0000000000..f8eaad3f86 --- /dev/null +++ b/tests/output/conftest.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026, 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 pathlib import Path + +import pytest + + +@pytest.fixture +def data_dir(test_root: Path) -> Path: + return Path(__file__).parent / "data" diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/areas.txt b/tests/output/data/20260810-1420eco-thermal_groups/about-the-study/areas.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/areas.txt rename to tests/output/data/20260810-1420eco-thermal_groups/about-the-study/areas.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/comments.txt b/tests/output/data/20260810-1420eco-thermal_groups/about-the-study/comments.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/comments.txt rename to tests/output/data/20260810-1420eco-thermal_groups/about-the-study/comments.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/links.txt b/tests/output/data/20260810-1420eco-thermal_groups/about-the-study/links.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/links.txt rename to tests/output/data/20260810-1420eco-thermal_groups/about-the-study/links.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/parameters.ini b/tests/output/data/20260810-1420eco-thermal_groups/about-the-study/parameters.ini similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/parameters.ini rename to tests/output/data/20260810-1420eco-thermal_groups/about-the-study/parameters.ini diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/study.ini b/tests/output/data/20260810-1420eco-thermal_groups/about-the-study/study.ini similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/about-the-study/study.ini rename to tests/output/data/20260810-1420eco-thermal_groups/about-the-study/study.ini diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/annualSystemCost.txt b/tests/output/data/20260810-1420eco-thermal_groups/annualSystemCost.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/annualSystemCost.txt rename to tests/output/data/20260810-1420eco-thermal_groups/annualSystemCost.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/checkIntegrity.txt b/tests/output/data/20260810-1420eco-thermal_groups/checkIntegrity.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/checkIntegrity.txt rename to tests/output/data/20260810-1420eco-thermal_groups/checkIntegrity.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/@ all areas/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/@ all areas/values-monthly.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/@ all areas/values-monthly.txt rename to tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/@ all areas/values-monthly.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/es/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/es/values-monthly.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/es/values-monthly.txt rename to tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/es/values-monthly.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/fr/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/fr/values-monthly.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/fr/values-monthly.txt rename to tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00001/areas/fr/values-monthly.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/@ all areas/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/@ all areas/values-monthly.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/@ all areas/values-monthly.txt rename to tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/@ all areas/values-monthly.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/es/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/es/values-monthly.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/es/values-monthly.txt rename to tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/es/values-monthly.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/fr/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/fr/values-monthly.txt similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/fr/values-monthly.txt rename to tests/output/data/20260810-1420eco-thermal_groups/economy/mc-ind/00002/areas/fr/values-monthly.txt diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/execution_info.ini b/tests/output/data/20260810-1420eco-thermal_groups/execution_info.ini similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/execution_info.ini rename to tests/output/data/20260810-1420eco-thermal_groups/execution_info.ini diff --git a/tests/output/filestudy/data/20260810-1420eco-thermal_groups/info.antares-output b/tests/output/data/20260810-1420eco-thermal_groups/info.antares-output similarity index 100% rename from tests/output/filestudy/data/20260810-1420eco-thermal_groups/info.antares-output rename to tests/output/data/20260810-1420eco-thermal_groups/info.antares-output diff --git a/tests/output/filestudy/test_matrix_aggregattion_result.py b/tests/output/filestudy/test_matrix_aggregation_result.py similarity index 87% rename from tests/output/filestudy/test_matrix_aggregattion_result.py rename to tests/output/filestudy/test_matrix_aggregation_result.py index e79113a7d4..2d7f92037e 100644 --- a/tests/output/filestudy/test_matrix_aggregattion_result.py +++ b/tests/output/filestudy/test_matrix_aggregation_result.py @@ -19,17 +19,15 @@ @pytest.fixture -def data_dir() -> Path: - return Path(__file__).parent / "data" +def output_dir(data_dir: Path) -> Path: + return data_dir / "20260810-1420eco-thermal_groups" -def test_build_aggregates__different_thermal_groups(data_dir: Path) -> None: +def test_build_aggregates__different_thermal_groups(output_dir: Path) -> None: """ Checks that when areas have different variables, we only get the relevant ones for each area. Here areas fr and es have different thermal groups. """ - output_dir = data_dir / "20260810-1420eco-thermal_groups" - download = StudyDownloadDTO( type=StudyDownloadType.AREA, years=[1], @@ -63,8 +61,7 @@ def test_build_aggregates__different_thermal_groups(data_dir: Path) -> None: ] -def test_build_aggregates__district(data_dir: Path) -> None: - output_dir = data_dir / "20260810-1420eco-thermal_groups" +def test_build_aggregates__district(output_dir: Path) -> None: download = StudyDownloadDTO( type=StudyDownloadType.DISTRICT, diff --git a/tests/output/storage/v2/test_variables_parsing.py b/tests/output/storage/v2/test_variables_parsing.py new file mode 100644 index 0000000000..e2fda4e050 --- /dev/null +++ b/tests/output/storage/v2/test_variables_parsing.py @@ -0,0 +1,51 @@ +# Copyright (c) 2026, 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 pathlib import Path + +import pytest + +from antarest.output.filestudy.model import FileOutput, VariableDescription +from antarest.output.storage.v2.variables_parsing import parse_output_variables + + +@pytest.fixture +def output_dir(data_dir: Path) -> Path: + return data_dir / "20260810-1420eco-thermal_groups" + + +def test_extract_parsing_to_db(output_dir: Path) -> None: + output = FileOutput(output_dir) + parsing_result = parse_output_variables(output) + + # Notes: weird stuff in input data: CO2 emissions in MWh, and different naming for thermal production groups + assert parsing_result.mc_ind_areas.variables == [ + VariableDescription(name="CO2 EMIS.", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type=None), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None), + ] + assert parsing_result.mc_ind_areas.area_vars == { + "@ all areas": [0, 1, 2, 3, 4, 5, 6, 7, 8], + "es": [9, 10, 1, 2, 3, 4, 6, 5], # we correctly get ES_NUCLEAR as 10 + "fr": [9, 11, 1, 2, 3, 4, 6, 5], # we correctly get FR_NUCLEAR as 11 + } + + assert parsing_result.mc_all_areas.variables == [] + assert parsing_result.mc_all_areas.area_vars == {} From 2a355d99f0b1a7c2cb423851244c9abe3d534444 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 14:40:47 +0200 Subject: [PATCH 03/29] fix mc-all parsing Signed-off-by: Sylvain Leclerc --- antarest/output/filestudy/model.py | 2 +- .../storage/v2/test_variables_parsing.py | 183 +++++++++++++++++- 2 files changed, 183 insertions(+), 2 deletions(-) diff --git a/antarest/output/filestudy/model.py b/antarest/output/filestudy/model.py index e53090a910..190428027d 100644 --- a/antarest/output/filestudy/model.py +++ b/antarest/output/filestudy/model.py @@ -246,7 +246,7 @@ def get_mc_all_file( """ Returns the path corresponding to the specified data, if it exists. """ - element_type = "areas" if isinstance(file_type, MCIndAreasQueryFile) else "links" + element_type = "areas" if isinstance(file_type, MCAllAreasQueryFile) else "links" file_path = self.mc_all_dir / element_type / area_id / f"{file_type}-{frequency}.txt" return file_path if file_path.exists() else None diff --git a/tests/output/storage/v2/test_variables_parsing.py b/tests/output/storage/v2/test_variables_parsing.py index e2fda4e050..6f3eae7c20 100644 --- a/tests/output/storage/v2/test_variables_parsing.py +++ b/tests/output/storage/v2/test_variables_parsing.py @@ -22,7 +22,188 @@ def output_dir(data_dir: Path) -> Path: return data_dir / "20260810-1420eco-thermal_groups" -def test_extract_parsing_to_db(output_dir: Path) -> None: +def test_extract_area_variables(output_dir: Path) -> None: + output = FileOutput(output_dir) + parsing_result = parse_output_variables(output) + + # Notes: weird stuff in input data: CO2 emissions in MWh, and different naming for thermal production groups + assert parsing_result.mc_ind_areas.variables == [ + VariableDescription(name="CO2 EMIS.", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type=None), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None), + ] + assert parsing_result.mc_ind_areas.area_vars == { + "@ all areas": [0, 1, 2, 3, 4, 5, 6, 7, 8], + "es": [9, 10, 1, 2, 3, 4, 6, 5], # we correctly get ES_NUCLEAR as 10 + "fr": [9, 11, 1, 2, 3, 4, 6, 5], # we correctly get FR_NUCLEAR as 11 + } + + assert parsing_result.mc_all_areas.variables == [ + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="EXP"), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="std"), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="min"), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="max"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="EXP"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="std"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="min"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="max"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="EXP"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="std"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="min"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="max"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="EXP"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="std"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="min"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="max"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="EXP"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="std"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="min"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="max"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="EXP"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="std"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="min"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="max"), + VariableDescription(name="NODU", unit=None, statistic_type="EXP"), + VariableDescription(name="NODU", unit=None, statistic_type="std"), + VariableDescription(name="NODU", unit=None, statistic_type="min"), + VariableDescription(name="NODU", unit=None, statistic_type="max"), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type="EXP"), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type="std"), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type="min"), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type="max"), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type="EXP"), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type="std"), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type="min"), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type="max"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="EXP"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="std"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="min"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="max"), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type="EXP"), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type="std"), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type="min"), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type="max"), + ] + assert parsing_result.mc_all_areas.area_vars == { + "@ all areas": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + ], + "es": [ + 0, + 1, + 2, + 3, + 36, + 37, + 38, + 39, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 24, + 25, + 26, + 27, + 20, + 21, + 22, + 23, + ], + "fr": [ + 0, + 1, + 2, + 3, + 40, + 41, + 42, + 43, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 24, + 25, + 26, + 27, + 20, + 21, + 22, + 23, + ], + } + + +def test_extract_area_variables_to_db(output_dir: Path) -> None: output = FileOutput(output_dir) parsing_result = parse_output_variables(output) From 46375b080a9323b382f479fccfa4c7257e84c601 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 14:47:25 +0200 Subject: [PATCH 04/29] make parsing fetch data from the right dir Signed-off-by: Sylvain Leclerc --- antarest/output/filestudy/model.py | 24 +++++++++++++++---- .../output/storage/v2/variables_parsing.py | 10 ++++++-- 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/antarest/output/filestudy/model.py b/antarest/output/filestudy/model.py index 190428027d..4322f3b9bd 100644 --- a/antarest/output/filestudy/model.py +++ b/antarest/output/filestudy/model.py @@ -212,19 +212,35 @@ def get_mc_year_dir(self, year: int) -> Path: return self.mc_ind_dir / f"{year:05d}" @cached_property - def link_ids(self) -> tuple[str, ...]: + def mc_ind_link_ids(self) -> tuple[str, ...]: """ - IDs of links that have data in the output directory, sorted. + IDs of links that have data in mc-ind, sorted. """ return tuple(sorted(d.name for d in self.iter_links_dir(self.first_mc_year))) @cached_property - def area_ids(self) -> tuple[str, ...]: + def mc_ind_area_ids(self) -> tuple[str, ...]: """ - IDs of areas that have data in the output directory, sorted. + IDs of areas that have data in mc-ind, sorted. """ return tuple(sorted(d.name for d in self.iter_areas_dir(self.first_mc_year))) + @cached_property + def mc_all_link_ids(self) -> tuple[str, ...]: + """ + IDs of links that have data in mc-all, sorted. + """ + links_dir = self.mc_all_dir / "links" + return tuple(sorted(d.name for d in links_dir.iterdir())) + + @cached_property + def mc_all_area_ids(self) -> tuple[str, ...]: + """ + IDs of areas that have data in mc-all, sorted. + """ + areas_dir = self.mc_all_dir / "areas" + return tuple(sorted(d.name for d in areas_dir.iterdir())) + def iter_areas_dir(self, mc_year: int) -> Iterable[Path]: """ No ordering guarantee. diff --git a/antarest/output/storage/v2/variables_parsing.py b/antarest/output/storage/v2/variables_parsing.py index 874da9558d..54e4bf07d1 100644 --- a/antarest/output/storage/v2/variables_parsing.py +++ b/antarest/output/storage/v2/variables_parsing.py @@ -16,7 +16,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Callable +from typing import Callable, Iterable from sqlalchemy.orm import Session @@ -55,7 +55,9 @@ def parse_area_variables(file_output: FileOutput, aggregation: ScenarioAggregati vars: list[VariableDescription] = [] area_cols: dict[str, list[int]] = {} + # data source depends on aggregation type get_file: Callable[[str, MatrixFrequency], Path | None] + area_ids: Iterable[str] match aggregation: case "mc-ind": @@ -63,12 +65,16 @@ def get_file(element_id: str, freq: MatrixFrequency) -> Path | None: return file_output.get_mc_ind_file( file_output.first_mc_year, MCIndAreasQueryFile.VALUES, element_id, freq ) + + area_ids = file_output.mc_ind_area_ids case "mc-all": def get_file(element_id: str, freq: MatrixFrequency) -> Path | None: return file_output.get_mc_all_file(MCAllAreasQueryFile.VALUES, element_id, freq) - for element_id in file_output.area_ids: + area_ids = file_output.mc_all_area_ids + + for element_id in area_ids: # searching for the first existing "frequency" for freq in MatrixFrequency: if data_file := get_file(element_id, freq): From 0f46f4994c9a77d6393907085e51ac4cc62895d4 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 14:49:11 +0200 Subject: [PATCH 05/29] remove test, add mc-all files Signed-off-by: Sylvain Leclerc --- .../mc-all/areas/@ all areas/id-monthly.txt | 19 ++++++++++++ .../areas/@ all areas/values-monthly.txt | 19 ++++++++++++ .../mc-all/areas/es/details-monthly.txt | 19 ++++++++++++ .../economy/mc-all/areas/es/id-monthly.txt | 19 ++++++++++++ .../mc-all/areas/es/values-monthly.txt | 19 ++++++++++++ .../mc-all/areas/fr/details-monthly.txt | 19 ++++++++++++ .../economy/mc-all/areas/fr/id-monthly.txt | 19 ++++++++++++ .../mc-all/areas/fr/values-monthly.txt | 19 ++++++++++++ .../economy/mc-all/grid/areas.txt | 3 ++ .../economy/mc-all/grid/digest.txt | 21 ++++++++++++++ .../economy/mc-all/grid/links.txt | 1 + .../economy/mc-all/grid/thermal.txt | 3 ++ .../storage/v2/test_variables_parsing.py | 29 ------------------- 13 files changed, 180 insertions(+), 29 deletions(-) create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/id-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/values-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/details-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/id-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/values-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/details-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/id-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/values-monthly.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/areas.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/digest.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/links.txt create mode 100644 tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/thermal.txt diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/id-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/id-monthly.txt new file mode 100644 index 0000000000..11c019870b --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/id-monthly.txt @@ -0,0 +1,19 @@ +system set of areas id monthly + VARIABLES BEGIN END + 14 1 12 + +system monthly CO2 EMIS. CO2 EMIS. AVL DTG AVL DTG DTG MRG DTG MRG MAX MRG MAX MRG NP COST NP COST RES LOAD RES LOAD NODU NODU + Tons Tons MWh MWh MWh MWh MWh MWh Euro Euro MWh MWh + index month min max min max min max min max min max min max min max + 1 JAN 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 2 FEB 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 3 MAR 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 4 APR 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 5 MAY 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 6 JUN 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 7 JUL 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 8 AUG 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 9 SEP 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 10 OCT 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 11 NOV 1 1 1 1 2 1 1 1 1 1 1 1 1 1 + 12 DEC 1 1 1 1 2 1 1 1 1 1 1 1 1 1 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/values-monthly.txt new file mode 100644 index 0000000000..49ca848289 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/@ all areas/values-monthly.txt @@ -0,0 +1,19 @@ +system set of areas va monthly + VARIABLES BEGIN END + 36 1 12 + +system monthly CO2 EMIS. CO2 EMIS. CO2 EMIS. CO2 EMIS. AVL DTG AVL DTG AVL DTG AVL DTG DTG MRG DTG MRG DTG MRG DTG MRG MAX MRG MAX MRG MAX MRG MAX MRG NP COST NP COST NP COST NP COST RES LOAD RES LOAD RES LOAD RES LOAD NODU NODU NODU NODU ES_NUCLEAR_TH_PROD ES_NUCLEAR_TH_PROD ES_NUCLEAR_TH_PROD ES_NUCLEAR_TH_PROD FR_NUCLEAR_TH_PROD FR_NUCLEAR_TH_PROD FR_NUCLEAR_TH_PROD FR_NUCLEAR_TH_PROD + Tons Tons Tons Tons MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh Euro Euro Euro Euro MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh + index month EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max + 1 JAN 0 0 0 0 1488000 0 1488000 1488000 74400 74400 0 148800 520800 0 520800 520800 0 0 0 0 967200 0 967200 967200 1488 0 1488 1488 744000 0 744000 744000 669600 74400 595200 744000 + 2 FEB 0 0 0 0 1344000 0 1344000 1344000 67200 67200 0 134400 470400 0 470400 470400 0 0 0 0 873600 0 873600 873600 1344 0 1344 1344 672000 0 672000 672000 604800 67200 537600 672000 + 3 MAR 0 0 0 0 1488000 0 1488000 1488000 74400 74400 0 148800 520800 0 520800 520800 0 0 0 0 967200 0 967200 967200 1488 0 1488 1488 744000 0 744000 744000 669600 74400 595200 744000 + 4 APR 0 0 0 0 1440000 0 1440000 1440000 72000 72000 0 144000 504000 0 504000 504000 0 0 0 0 936000 0 936000 936000 1440 0 1440 1440 720000 0 720000 720000 648000 72000 576000 720000 + 5 MAY 0 0 0 0 1488000 0 1488000 1488000 74400 74400 0 148800 520800 0 520800 520800 0 0 0 0 967200 0 967200 967200 1488 0 1488 1488 744000 0 744000 744000 669600 74400 595200 744000 + 6 JUN 0 0 0 0 1440000 0 1440000 1440000 72000 72000 0 144000 504000 0 504000 504000 0 0 0 0 936000 0 936000 936000 1440 0 1440 1440 720000 0 720000 720000 648000 72000 576000 720000 + 7 JUL 0 0 0 0 1488000 0 1488000 1488000 74400 74400 0 148800 520800 0 520800 520800 0 0 0 0 967200 0 967200 967200 1488 0 1488 1488 744000 0 744000 744000 669600 74400 595200 744000 + 8 AUG 0 0 0 0 1488000 0 1488000 1488000 74400 74400 0 148800 520800 0 520800 520800 0 0 0 0 967200 0 967200 967200 1488 0 1488 1488 744000 0 744000 744000 669600 74400 595200 744000 + 9 SEP 0 0 0 0 1440000 0 1440000 1440000 72000 72000 0 144000 504000 0 504000 504000 0 0 0 0 936000 0 936000 936000 1440 0 1440 1440 720000 0 720000 720000 648000 72000 576000 720000 + 10 OCT 0 0 0 0 1488000 0 1488000 1488000 74400 74400 0 148800 520800 0 520800 520800 0 0 0 0 967200 0 967200 967200 1488 0 1488 1488 744000 0 744000 744000 669600 74400 595200 744000 + 11 NOV 0 0 0 0 1440000 0 1440000 1440000 72000 72000 0 144000 504000 0 504000 504000 0 0 0 0 936000 0 936000 936000 1440 0 1440 1440 720000 0 720000 720000 648000 72000 576000 720000 + 12 DEC 0 0 0 0 1440000 0 1440000 1440000 72000 72000 0 144000 504000 0 504000 504000 0 0 0 0 936000 0 936000 936000 1440 0 1440 1440 720000 0 720000 720000 648000 72000 576000 720000 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/details-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/details-monthly.txt new file mode 100644 index 0000000000..65b2d97ce4 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/details-monthly.txt @@ -0,0 +1,19 @@ +es area de monthly + VARIABLES BEGIN END + 5 1 12 + +es monthly es_nuc es_nuc es_nuc es_nuc es_nuc + MWh MIN GEN - MWh NP Cost - Euro NODU Profit - Euro + index month EXP EXP EXP EXP EXP + 1 JAN 744000 0 0 744 -37200440 + 2 FEB 672000 0 0 672 -33600397 + 3 MAR 744000 0 0 744 -37200440 + 4 APR 720000 0 0 720 -36000426 + 5 MAY 744000 0 0 744 -37200440 + 6 JUN 720000 0 0 720 -36000426 + 7 JUL 744000 0 0 744 -37200440 + 8 AUG 744000 0 0 744 -37200440 + 9 SEP 720000 0 0 720 -36000426 + 10 OCT 744000 0 0 744 -37200440 + 11 NOV 720000 0 0 720 -36000426 + 12 DEC 720000 0 0 720 -36000426 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/id-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/id-monthly.txt new file mode 100644 index 0000000000..16f32d1a0e --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/id-monthly.txt @@ -0,0 +1,19 @@ +es area id monthly + VARIABLES BEGIN END + 16 1 12 + +es monthly CO2 EMIS. CO2 EMIS. ES_NUCLEAR ES_NUCLEAR AVL DTG AVL DTG DTG MRG DTG MRG MAX MRG MAX MRG NP COST NP COST NODU NODU RES LOAD RES LOAD + Tons Tons MWh MWh MWh MWh MWh MWh MWh MWh Euro Euro MWh MWh + index month min max min max min max min max min max min max min max min max + 1 JAN 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 2 FEB 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 3 MAR 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 4 APR 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 5 MAY 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 6 JUN 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 7 JUL 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 8 AUG 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 9 SEP 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 10 OCT 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 11 NOV 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 + 12 DEC 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/values-monthly.txt new file mode 100644 index 0000000000..b27d7c1661 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/es/values-monthly.txt @@ -0,0 +1,19 @@ +es area va monthly + VARIABLES BEGIN END + 32 1 12 + +es monthly CO2 EMIS. CO2 EMIS. CO2 EMIS. CO2 EMIS. ES_NUCLEAR ES_NUCLEAR ES_NUCLEAR ES_NUCLEAR AVL DTG AVL DTG AVL DTG AVL DTG DTG MRG DTG MRG DTG MRG DTG MRG MAX MRG MAX MRG MAX MRG MAX MRG NP COST NP COST NP COST NP COST NODU NODU NODU NODU RES LOAD RES LOAD RES LOAD RES LOAD + Tons Tons Tons Tons MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh Euro Euro Euro Euro MWh MWh MWh MWh + index month EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max + 1 JAN 0 0 0 0 744000 0 744000 744000 744000 0 744000 744000 0 0 0 0 372000 0 372000 372000 0 0 0 0 744 0 744 744 372000 0 372000 372000 + 2 FEB 0 0 0 0 672000 0 672000 672000 672000 0 672000 672000 0 0 0 0 336000 0 336000 336000 0 0 0 0 672 0 672 672 336000 0 336000 336000 + 3 MAR 0 0 0 0 744000 0 744000 744000 744000 0 744000 744000 0 0 0 0 372000 0 372000 372000 0 0 0 0 744 0 744 744 372000 0 372000 372000 + 4 APR 0 0 0 0 720000 0 720000 720000 720000 0 720000 720000 0 0 0 0 360000 0 360000 360000 0 0 0 0 720 0 720 720 360000 0 360000 360000 + 5 MAY 0 0 0 0 744000 0 744000 744000 744000 0 744000 744000 0 0 0 0 372000 0 372000 372000 0 0 0 0 744 0 744 744 372000 0 372000 372000 + 6 JUN 0 0 0 0 720000 0 720000 720000 720000 0 720000 720000 0 0 0 0 360000 0 360000 360000 0 0 0 0 720 0 720 720 360000 0 360000 360000 + 7 JUL 0 0 0 0 744000 0 744000 744000 744000 0 744000 744000 0 0 0 0 372000 0 372000 372000 0 0 0 0 744 0 744 744 372000 0 372000 372000 + 8 AUG 0 0 0 0 744000 0 744000 744000 744000 0 744000 744000 0 0 0 0 372000 0 372000 372000 0 0 0 0 744 0 744 744 372000 0 372000 372000 + 9 SEP 0 0 0 0 720000 0 720000 720000 720000 0 720000 720000 0 0 0 0 360000 0 360000 360000 0 0 0 0 720 0 720 720 360000 0 360000 360000 + 10 OCT 0 0 0 0 744000 0 744000 744000 744000 0 744000 744000 0 0 0 0 372000 0 372000 372000 0 0 0 0 744 0 744 744 372000 0 372000 372000 + 11 NOV 0 0 0 0 720000 0 720000 720000 720000 0 720000 720000 0 0 0 0 360000 0 360000 360000 0 0 0 0 720 0 720 720 360000 0 360000 360000 + 12 DEC 0 0 0 0 720000 0 720000 720000 720000 0 720000 720000 0 0 0 0 360000 0 360000 360000 0 0 0 0 720 0 720 720 360000 0 360000 360000 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/details-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/details-monthly.txt new file mode 100644 index 0000000000..b83dcf7212 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/details-monthly.txt @@ -0,0 +1,19 @@ +fr area de monthly + VARIABLES BEGIN END + 5 1 12 + +fr monthly fr_nuclear fr_nuclear fr_nuclear fr_nuclear fr_nuclear + MWh MIN GEN - MWh NP Cost - Euro NODU Profit - Euro + index month EXP EXP EXP EXP EXP + 1 JAN 669600 0 0 744 -347 + 2 FEB 604800 0 0 672 -313 + 3 MAR 669600 0 0 744 -347 + 4 APR 648000 0 0 720 -336 + 5 MAY 669600 0 0 744 -347 + 6 JUN 648000 0 0 720 -336 + 7 JUL 669600 0 0 744 -347 + 8 AUG 669600 0 0 744 -347 + 9 SEP 648000 0 0 720 -336 + 10 OCT 669600 0 0 744 -347 + 11 NOV 648000 0 0 720 -336 + 12 DEC 648000 0 0 720 -336 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/id-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/id-monthly.txt new file mode 100644 index 0000000000..8a6828012f --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/id-monthly.txt @@ -0,0 +1,19 @@ +fr area id monthly + VARIABLES BEGIN END + 16 1 12 + +fr monthly CO2 EMIS. CO2 EMIS. FR_NUCLEAR FR_NUCLEAR AVL DTG AVL DTG DTG MRG DTG MRG MAX MRG MAX MRG NP COST NP COST NODU NODU RES LOAD RES LOAD + Tons Tons MWh MWh MWh MWh MWh MWh MWh MWh Euro Euro MWh MWh + index month min max min max min max min max min max min max min max min max + 1 JAN 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 2 FEB 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 3 MAR 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 4 APR 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 5 MAY 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 6 JUN 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 7 JUL 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 8 AUG 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 9 SEP 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 10 OCT 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 11 NOV 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 + 12 DEC 1 1 1 2 1 1 2 1 1 1 1 1 1 1 1 1 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/values-monthly.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/values-monthly.txt new file mode 100644 index 0000000000..aa75253bc5 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/areas/fr/values-monthly.txt @@ -0,0 +1,19 @@ +fr area va monthly + VARIABLES BEGIN END + 32 1 12 + +fr monthly CO2 EMIS. CO2 EMIS. CO2 EMIS. CO2 EMIS. FR_NUCLEAR FR_NUCLEAR FR_NUCLEAR FR_NUCLEAR AVL DTG AVL DTG AVL DTG AVL DTG DTG MRG DTG MRG DTG MRG DTG MRG MAX MRG MAX MRG MAX MRG MAX MRG NP COST NP COST NP COST NP COST NODU NODU NODU NODU RES LOAD RES LOAD RES LOAD RES LOAD + Tons Tons Tons Tons MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh MWh Euro Euro Euro Euro MWh MWh MWh MWh + index month EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max EXP std min max + 1 JAN 0 0 0 0 669600 74400 595200 744000 744000 0 744000 744000 74400 74400 0 148800 148800 0 148800 148800 0 0 0 0 744 0 744 744 595200 0 595200 595200 + 2 FEB 0 0 0 0 604800 67200 537600 672000 672000 0 672000 672000 67200 67200 0 134400 134400 0 134400 134400 0 0 0 0 672 0 672 672 537600 0 537600 537600 + 3 MAR 0 0 0 0 669600 74400 595200 744000 744000 0 744000 744000 74400 74400 0 148800 148800 0 148800 148800 0 0 0 0 744 0 744 744 595200 0 595200 595200 + 4 APR 0 0 0 0 648000 72000 576000 720000 720000 0 720000 720000 72000 72000 0 144000 144000 0 144000 144000 0 0 0 0 720 0 720 720 576000 0 576000 576000 + 5 MAY 0 0 0 0 669600 74400 595200 744000 744000 0 744000 744000 74400 74400 0 148800 148800 0 148800 148800 0 0 0 0 744 0 744 744 595200 0 595200 595200 + 6 JUN 0 0 0 0 648000 72000 576000 720000 720000 0 720000 720000 72000 72000 0 144000 144000 0 144000 144000 0 0 0 0 720 0 720 720 576000 0 576000 576000 + 7 JUL 0 0 0 0 669600 74400 595200 744000 744000 0 744000 744000 74400 74400 0 148800 148800 0 148800 148800 0 0 0 0 744 0 744 744 595200 0 595200 595200 + 8 AUG 0 0 0 0 669600 74400 595200 744000 744000 0 744000 744000 74400 74400 0 148800 148800 0 148800 148800 0 0 0 0 744 0 744 744 595200 0 595200 595200 + 9 SEP 0 0 0 0 648000 72000 576000 720000 720000 0 720000 720000 72000 72000 0 144000 144000 0 144000 144000 0 0 0 0 720 0 720 720 576000 0 576000 576000 + 10 OCT 0 0 0 0 669600 74400 595200 744000 744000 0 744000 744000 74400 74400 0 148800 148800 0 148800 148800 0 0 0 0 744 0 744 744 595200 0 595200 595200 + 11 NOV 0 0 0 0 648000 72000 576000 720000 720000 0 720000 720000 72000 72000 0 144000 144000 0 144000 144000 0 0 0 0 720 0 720 720 576000 0 576000 576000 + 12 DEC 0 0 0 0 648000 72000 576000 720000 720000 0 720000 720000 72000 72000 0 144000 144000 0 144000 144000 0 0 0 0 720 0 720 720 576000 0 576000 576000 diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/areas.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/areas.txt new file mode 100644 index 0000000000..cd30762d89 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/areas.txt @@ -0,0 +1,3 @@ +id name +es es +fr fr diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/digest.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/digest.txt new file mode 100644 index 0000000000..4b8e6cab23 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/digest.txt @@ -0,0 +1,21 @@ + digest + VARIABLES AREAS LINKS + 7 2 0 + + CO2 EMIS. AVL DTG DTG MRG MAX MRG NP COST NODU RES LOAD + Tons MWh MWh MWh Euro MWh + EXP EXP EXP EXP EXP EXP EXP + es 0 8736000 0 4368000 0 8736 4368000 + fr 0 8736000 873600 1747200 0 8736 6988800 + + + digest + VARIABLES AREAS LINKS + 7 1 0 + + CO2 EMIS. AVL DTG DTG MRG MAX MRG NP COST RES LOAD NODU + Tons MWh MWh MWh Euro MWh + EXP EXP EXP EXP EXP EXP EXP + @ All areas 0 17472000 873600 6115200 0 11356800 17472 + + diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/links.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/links.txt new file mode 100644 index 0000000000..9813cf61d1 --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/links.txt @@ -0,0 +1 @@ +upstream downstream diff --git a/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/thermal.txt b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/thermal.txt new file mode 100644 index 0000000000..f7c59d80fc --- /dev/null +++ b/tests/output/data/20260810-1420eco-thermal_groups/economy/mc-all/grid/thermal.txt @@ -0,0 +1,3 @@ +area id id name group unit count nominal capacity min stable power min up/down time spinning co2 marginal cost fixed cost startup cost market bid cost spread cost +es es_nuc es_nuc ES_NUCLEAR 1 1000.000000 0.000000 1 1 0.000000 0.000000 50.000000 0.000000 0.000000 0.000000 0.000000 +fr fr_nuclear fr_nuclear FR_NUCLEAR 1 1000.000000 0.000000 1 1 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 diff --git a/tests/output/storage/v2/test_variables_parsing.py b/tests/output/storage/v2/test_variables_parsing.py index 6f3eae7c20..16ba8d1cdc 100644 --- a/tests/output/storage/v2/test_variables_parsing.py +++ b/tests/output/storage/v2/test_variables_parsing.py @@ -201,32 +201,3 @@ def test_extract_area_variables(output_dir: Path) -> None: 23, ], } - - -def test_extract_area_variables_to_db(output_dir: Path) -> None: - output = FileOutput(output_dir) - parsing_result = parse_output_variables(output) - - # Notes: weird stuff in input data: CO2 emissions in MWh, and different naming for thermal production groups - assert parsing_result.mc_ind_areas.variables == [ - VariableDescription(name="CO2 EMIS.", unit="MWh", statistic_type=None), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), - VariableDescription(name="NP COST", unit="Euro", statistic_type=None), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), - VariableDescription(name="NODU", unit=None, statistic_type=None), - VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), - VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), - VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type=None), - VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None), - ] - assert parsing_result.mc_ind_areas.area_vars == { - "@ all areas": [0, 1, 2, 3, 4, 5, 6, 7, 8], - "es": [9, 10, 1, 2, 3, 4, 6, 5], # we correctly get ES_NUCLEAR as 10 - "fr": [9, 11, 1, 2, 3, 4, 6, 5], # we correctly get FR_NUCLEAR as 11 - } - - assert parsing_result.mc_all_areas.variables == [] - assert parsing_result.mc_all_areas.area_vars == {} From fd259be9a3cbd68b26ac767c3df49480b9f4af31 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 14:54:19 +0200 Subject: [PATCH 06/29] move functions from dbmodel Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/dbmodel.py | 39 +------------ .../output/storage/v2/variables_fetching.py | 55 +++++++++++++++++++ 2 files changed, 58 insertions(+), 36 deletions(-) create mode 100644 antarest/output/storage/v2/variables_fetching.py diff --git a/antarest/output/storage/v2/dbmodel.py b/antarest/output/storage/v2/dbmodel.py index 0d0f0936e3..b84cc48104 100644 --- a/antarest/output/storage/v2/dbmodel.py +++ b/antarest/output/storage/v2/dbmodel.py @@ -9,14 +9,13 @@ # SPDX-License-Identifier: MPL-2.0 # # This file is part of the Antares project. -from typing import Any, Iterable, Literal, TypeAlias +from typing import Any, Literal, TypeAlias -from sqlalchemy import BigInteger, Dialect, ForeignKeyConstraint, SmallInteger, String, select, types -from sqlalchemy.orm import Mapped, Session, mapped_column +from sqlalchemy import BigInteger, Dialect, ForeignKeyConstraint, SmallInteger, String, types +from sqlalchemy.orm import Mapped, mapped_column from typing_extensions import override from antarest.dbmodel import Base -from antarest.output.filestudy.model import VariableDescription ElementType: TypeAlias = Literal[ "area", @@ -84,35 +83,3 @@ class DbParquetArea(Base): area_id: Mapped[str] = mapped_column(primary_key=True) mc_all_vars: Mapped[list[int]] = mapped_column(Columns) mc_ind_vars: Mapped[list[int]] = mapped_column(Columns) - - -class VariablesIndex: - def __init__(self, variables: Iterable[DbParquetVariable]) -> None: - self._variables: dict[tuple[ScenarioAggregation, ElementType], list[VariableDescription]] = {} - for v in variables: - self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(to_var_model(v)) - - def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: - return self._variables.get((aggregation, element_type), []) - - -def to_var_model(db_var: DbParquetVariable) -> VariableDescription: - return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) - - -def get_area_variables( - session: Session, output_id: int, aggregation: ScenarioAggregation, area_id: str -) -> list[VariableDescription]: - - # All variables, should load fast ? - output_variables = session.execute( - select(DbParquetVariable).where(DbParquetVariable.output_id == output_id) - ).scalars() - variables_index = VariablesIndex(output_variables) - - # Get area information - area = session.execute(select(DbParquetArea).where(DbParquetArea.area_id == area_id)).scalar_one() - - all_areas_vars = variables_index.get_variables(aggregation, "area") - cols = area.mc_all_vars if aggregation == "mc-all" else area.mc_ind_vars - return [all_areas_vars[c] for c in cols] diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py new file mode 100644 index 0000000000..af650ae646 --- /dev/null +++ b/antarest/output/storage/v2/variables_fetching.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026, 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. + +""" +Retrieving variable-related information from the database +""" + +from typing import Iterable + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from antarest.output.filestudy.model import VariableDescription +from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ElementType, ScenarioAggregation + + +class VariablesIndex: + def __init__(self, variables: Iterable[DbParquetVariable]) -> None: + self._variables: dict[tuple[ScenarioAggregation, ElementType], list[VariableDescription]] = {} + for v in variables: + self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(to_var_model(v)) + + def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: + return self._variables.get((aggregation, element_type), []) + + +def to_var_model(db_var: DbParquetVariable) -> VariableDescription: + return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) + + +def get_area_variables( + session: Session, output_id: int, aggregation: ScenarioAggregation, area_id: str +) -> list[VariableDescription]: + + # All variables, should load fast ? + output_variables = session.execute( + select(DbParquetVariable).where(DbParquetVariable.output_id == output_id) + ).scalars() + variables_index = VariablesIndex(output_variables) + + # Get area information + area = session.execute(select(DbParquetArea).where(DbParquetArea.area_id == area_id)).scalar_one() + + all_areas_vars = variables_index.get_variables(aggregation, "area") + cols = area.mc_all_vars if aggregation == "mc-all" else area.mc_ind_vars + return [all_areas_vars[c] for c in cols] From 2b5e48da5e3676ceaf443e9cad99433c8ab1bfa8 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 15:08:09 +0200 Subject: [PATCH 07/29] unit test for fetch from database Signed-off-by: Sylvain Leclerc --- .../output/storage/v2/variables_fetching.py | 11 ++- .../storage/v2/test_variables_fetching.py | 93 +++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) create mode 100644 tests/output/storage/v2/test_variables_fetching.py diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index af650ae646..85c3393eb3 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -24,16 +24,23 @@ class VariablesIndex: + """ + Helper class to retrieve variables from DB models. + """ + def __init__(self, variables: Iterable[DbParquetVariable]) -> None: self._variables: dict[tuple[ScenarioAggregation, ElementType], list[VariableDescription]] = {} for v in variables: - self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(to_var_model(v)) + self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(_to_var_model(v)) def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: + """ + Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) + """ return self._variables.get((aggregation, element_type), []) -def to_var_model(db_var: DbParquetVariable) -> VariableDescription: +def _to_var_model(db_var: DbParquetVariable) -> VariableDescription: return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) diff --git a/tests/output/storage/v2/test_variables_fetching.py b/tests/output/storage/v2/test_variables_fetching.py new file mode 100644 index 0000000000..2fbe8be389 --- /dev/null +++ b/tests/output/storage/v2/test_variables_fetching.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026, 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 pathlib import Path + +import pytest +from sqlalchemy.orm import Session + +from antarest.output.filestudy.model import FileOutput, VariableDescription +from antarest.output.storage.v2.dbmodel import DbParquetOutput +from antarest.output.storage.v2.variables_fetching import get_area_variables +from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database + + +@pytest.fixture +def output_dir(data_dir: Path) -> Path: + return data_dir / "20260810-1420eco-thermal_groups" + + +def test_get_area_variables(db_session: Session, output_dir: Path) -> None: + + db_output = DbParquetOutput(id=0) + db_session.add(db_output) + db_session.flush() + + file_output = FileOutput(output_dir) + extract_output_variables_to_database(db_session, db_output.id, file_output) + db_session.flush() + + assert get_area_variables(db_session, db_output.id, "mc-ind", "fr") == [ + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + ] + + assert get_area_variables(db_session, db_output.id, "mc-ind", "es") == [ + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + ] + + assert get_area_variables(db_session, db_output.id, "mc-all", "es") == [ + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="EXP"), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="std"), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="min"), + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="max"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="EXP"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="std"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="min"), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="max"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="EXP"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="std"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="min"), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type="max"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="EXP"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="std"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="min"), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type="max"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="EXP"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="std"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="min"), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type="max"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="EXP"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="std"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="min"), + VariableDescription(name="NP COST", unit="Euro", statistic_type="max"), + VariableDescription(name="NODU", unit=None, statistic_type="EXP"), + VariableDescription(name="NODU", unit=None, statistic_type="std"), + VariableDescription(name="NODU", unit=None, statistic_type="min"), + VariableDescription(name="NODU", unit=None, statistic_type="max"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="EXP"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="std"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="min"), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type="max"), + ] From 56c3ffb73aa4af3122c354bf6783e8a37dea0eff Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 17:41:00 +0200 Subject: [PATCH 08/29] add some todos, start re-working parquet extraction Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/storage.py | 1 + .../output/storage/v2/variables_fetching.py | 43 ++++++++++++++++--- .../output/storage/v2/variables_parsing.py | 2 +- .../output/storage/v2/variables_storage.py | 40 +++++++++++++++++ 4 files changed, 79 insertions(+), 7 deletions(-) diff --git a/antarest/output/storage/v2/storage.py b/antarest/output/storage/v2/storage.py index c2364db47a..b2e8b098cc 100644 --- a/antarest/output/storage/v2/storage.py +++ b/antarest/output/storage/v2/storage.py @@ -222,6 +222,7 @@ def import_output( simulation_range = _extract_simulation_range(dir_path) + # TODO: first, extract variables metadata to database variables_target = parquet_output_dir(self._variables_dir, study_id, output_name) extract_output_to_parquet(dir_path, variables_target) diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index 85c3393eb3..fe37c013d3 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -11,11 +11,13 @@ # This file is part of the Antares project. """ -Retrieving variable-related information from the database +Fetching variables metadata from the database """ +from dataclasses import dataclass from typing import Iterable +from pyarrow.lib import Sequence from sqlalchemy import select from sqlalchemy.orm import Session @@ -23,24 +25,53 @@ from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ElementType, ScenarioAggregation +@dataclass(frozen=True) +class VariableColumn: + """ + Attributes: + column: offset of the column in the parquet file (actual col will be number of index cols + this offset) + """ + + column: int + name: str + unit: str | None + statistic_type: str | None + + class VariablesIndex: """ - Helper class to retrieve variables from DB models. + Helper class to retrieve variable info from DB models. """ def __init__(self, variables: Iterable[DbParquetVariable]) -> None: - self._variables: dict[tuple[ScenarioAggregation, ElementType], list[VariableDescription]] = {} + self._variables: dict[tuple[ScenarioAggregation, ElementType], list[DbParquetVariable]] = {} for v in variables: - self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(_to_var_model(v)) + self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(v) - def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: + def _get_db_vars(self, aggregation: ScenarioAggregation, element_type: ElementType) -> Sequence[DbParquetVariable]: """ Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) """ return self._variables.get((aggregation, element_type), []) + def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: + """ + Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) + """ + return [_to_var_desc(v) for v in self._get_db_vars(aggregation, element_type)] + + def get_variable_columns(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableColumn]: + """ + Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) + """ + return [_to_var_col(v) for v in self._get_db_vars(aggregation, element_type)] + + +def _to_var_col(db_var: DbParquetVariable) -> VariableColumn: + return VariableColumn(db_var.column, db_var.name, db_var.unit, db_var.statistic_type) + -def _to_var_model(db_var: DbParquetVariable) -> VariableDescription: +def _to_var_desc(db_var: DbParquetVariable) -> VariableDescription: return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) diff --git a/antarest/output/storage/v2/variables_parsing.py b/antarest/output/storage/v2/variables_parsing.py index 54e4bf07d1..e6ce0472bd 100644 --- a/antarest/output/storage/v2/variables_parsing.py +++ b/antarest/output/storage/v2/variables_parsing.py @@ -11,7 +11,7 @@ # This file is part of the Antares project. """ -Extraction of variables metadata from file studies, in order to populate the database +Parsing of variables metadata from file studies, in order to populate the database """ from dataclasses import dataclass diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 4b8e1dcb3a..6e1150bd77 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -23,11 +23,14 @@ import tempfile from collections.abc import Iterator, Sequence from pathlib import Path +from typing import Iterable import polars as pl +import pyarrow as pa from antarest.core.exceptions import MCRootNotHandled, OutputAggregationError, OutputNotFound, OutputSubFolderNotFound from antarest.core.serde.parquet_writer import ( + BatchParquetWriter, write_dataframes_in_parquet_format_by_column_sets, write_dataframes_stream_parquet, yield_dataframes_from_parquet, @@ -37,15 +40,19 @@ from antarest.output.filestudy.model import ( MCYEAR_COL, TIME_ID_COL, + FileOutput, MCAllAreasQueryFile, MCAllLinksQueryFile, MCIndAreasQueryFile, MCIndLinksQueryFile, MCRoot, + OutputDataFrame, QueryFileType, + VariableDescription, find_mode_dir, get_output_object_type, ) +from antarest.output.storage.v2.variables_fetching import VariablesIndex from antarest.study.model import MatrixFrequency logger = logging.getLogger(__name__) @@ -125,6 +132,39 @@ def _aggregate_to_parquet( _merge_intermediate_parquets(file_paths, new_index, target_path) +def _iter_mc_ind_area_files( + file_output: FileOutput, freq: MatrixFrequency +) -> Iterable[OutputDataFrame[VariableDescription]]: + start_col = get_start_column(freq) + for mc_year in file_output.mc_years: + for area_id in file_output.mc_ind_area_ids: + if data_file := file_output.get_mc_ind_file(mc_year, MCIndAreasQueryFile.VALUES, area_id, freq): + yield parse_output_file(data_file, start_col) + + +def _extract_areas( + index: VariablesIndex, + file_output: FileOutput, + target_dir: Path, +) -> None: + + variable_cols = index.get_variable_columns("mc-ind", "area") + + index_fields = [pa.field("area", pa.string()), pa.field("mcYear", pa.int32()), pa.field("timeId", pa.int32())] + schema = pa.schema(index_fields + [pa.field(str(i), pa.float64()) for i in range(len(variable_cols))]) + + for freq in MatrixFrequency: + output_file_path = target_dir / f"mc-ind_areas_{freq.value}.parquet" + with BatchParquetWriter(output_file_path, schema) as writer: + for df in _iter_mc_ind_area_files(file_output, freq): + data = df.data + # TODO: here we need to reshape the DF to comply with the schema (add index, add missing cols, + # reorder cols) + writer.add_table(data.to_arrow()) + + +# TODO: see above, this implementation needs to be replaced with one that uses the +# column indices that have been determined when parsing variable metadata def _extract_areas( output_dir: Path, base_path: Path, From fd371a7b6fbe3f21624ff8344ad817f728e49a49 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 17:42:33 +0200 Subject: [PATCH 09/29] remove unused fixture dependency Signed-off-by: Sylvain Leclerc --- tests/output/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/output/conftest.py b/tests/output/conftest.py index f8eaad3f86..c589615d0d 100644 --- a/tests/output/conftest.py +++ b/tests/output/conftest.py @@ -15,5 +15,5 @@ @pytest.fixture -def data_dir(test_root: Path) -> Path: +def data_dir() -> Path: return Path(__file__).parent / "data" From fdb781fa3f2fe72576470febfa9c1361aec31927 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Fri, 21 Aug 2026 17:58:44 +0200 Subject: [PATCH 10/29] add more comments Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/dbmodel.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/antarest/output/storage/v2/dbmodel.py b/antarest/output/storage/v2/dbmodel.py index b84cc48104..a1da0e94b0 100644 --- a/antarest/output/storage/v2/dbmodel.py +++ b/antarest/output/storage/v2/dbmodel.py @@ -31,6 +31,8 @@ class DbParquetOutput(Base): # TODO: we should merge the existing v2_output_metadata tables into this one + # the integer identifier will be easier and more efficient to use than the couple of strings + # study_id / output_id __tablename__ = "parquet_output" @@ -38,6 +40,16 @@ class DbParquetOutput(Base): class DbParquetVariable(Base): + """ + Represents one of the variables referenced in an output. + + Those variables are then referenced by elements of the system (areas, links ...), that contain + actual data for them. + + Attributes: + column: the column offset in the actual parquet file, compared to index columns (starts at 0). + """ + __tablename__ = "parquet_variable" __table_args__ = (ForeignKeyConstraint(["output_id"], ["parquet_output.id"]),) # TODO @@ -75,6 +87,13 @@ def process_result_value(self, value: Any | None, dialect: Dialect) -> list[int] class DbParquetArea(Base): + """ + Information related to an area of an output, in particular which variables it has data for, + in mc-ind and in mc-all (they may differ). + + The variables are reference through their column index. + """ + __tablename__ = "parquet_area" __table_args__ = (ForeignKeyConstraint(["output_id"], ["parquet_output.id"]),) # TODO @@ -83,3 +102,6 @@ class DbParquetArea(Base): area_id: Mapped[str] = mapped_column(primary_key=True) mc_all_vars: Mapped[list[int]] = mapped_column(Columns) mc_ind_vars: Mapped[list[int]] = mapped_column(Columns) + + +# TODO: add tables for other element types: links, thermal clusters, etc From 8c17bf0e17a9bef59efcb6974ecc8fa450cdfbc2 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Sat, 22 Aug 2026 10:39:55 +0200 Subject: [PATCH 11/29] implement columns adaptation Signed-off-by: Sylvain Leclerc --- .../output/storage/v2/variables_fetching.py | 9 ++-- .../output/storage/v2/variables_storage.py | 34 +++++++++++--- .../storage/v2/test_variable_storage.py | 46 +++++++++++++++++++ 3 files changed, 79 insertions(+), 10 deletions(-) create mode 100644 tests/output/storage/v2/test_variable_storage.py diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index fe37c013d3..5567ee2bc6 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -15,9 +15,8 @@ """ from dataclasses import dataclass -from typing import Iterable +from typing import Iterable, Sequence -from pyarrow.lib import Sequence from sqlalchemy import select from sqlalchemy.orm import Session @@ -44,9 +43,11 @@ class VariablesIndex: """ def __init__(self, variables: Iterable[DbParquetVariable]) -> None: - self._variables: dict[tuple[ScenarioAggregation, ElementType], list[DbParquetVariable]] = {} + vars: dict[tuple[ScenarioAggregation, ElementType], list[DbParquetVariable]] = {} for v in variables: - self._variables.setdefault((v.scenario_aggregation, v.element_type), []).append(v) + vars.setdefault((v.scenario_aggregation, v.element_type), []).append(v) + + self._variables = {k: sorted(v, key=lambda v: v.column) for k, v in vars.items()} # sort by columns def _get_db_vars(self, aggregation: ScenarioAggregation, element_type: ElementType) -> Sequence[DbParquetVariable]: """ diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 6e1150bd77..2d0f004a68 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -23,10 +23,13 @@ import tempfile from collections.abc import Iterator, Sequence from pathlib import Path -from typing import Iterable +from typing import Any, Iterable import polars as pl +import polars.selectors as pls import pyarrow as pa +from polars import Float64 +from pyarrow.lib import Field from antarest.core.exceptions import MCRootNotHandled, OutputAggregationError, OutputNotFound, OutputSubFolderNotFound from antarest.core.serde.parquet_writer import ( @@ -142,24 +145,43 @@ def _iter_mc_ind_area_files( yield parse_output_file(data_file, start_col) +def _adapt_df(variable_cols: list[VariableDescription], df: OutputDataFrame[VariableDescription]) -> pl.DataFrame: + """ + Reshapes the DF so that the column order matches the list of variables, and fills missing variables with null cols + + TODO: create a writer class which encapsulates this plus the writing ? + """ + col_for_variable = {v: i for i, v in enumerate(df.headers)} + nulls = pl.lit(None, dtype=Float64()) + res = df.data.select( + [ + pls.by_index(col_for_variable[v]).alias(str(i)) if v in col_for_variable else nulls.alias(str(i)) + for i, v in enumerate(variable_cols) + ] + ) + return res + + def _extract_areas( index: VariablesIndex, file_output: FileOutput, target_dir: Path, ) -> None: - variable_cols = index.get_variable_columns("mc-ind", "area") + variable_cols = index.get_variables("mc-ind", "area") - index_fields = [pa.field("area", pa.string()), pa.field("mcYear", pa.int32()), pa.field("timeId", pa.int32())] + index_fields: list[Field[Any]] = [ + pa.field("area", pa.string()), + pa.field("mcYear", pa.int32()), + pa.field("timeId", pa.int32()), + ] schema = pa.schema(index_fields + [pa.field(str(i), pa.float64()) for i in range(len(variable_cols))]) for freq in MatrixFrequency: output_file_path = target_dir / f"mc-ind_areas_{freq.value}.parquet" with BatchParquetWriter(output_file_path, schema) as writer: for df in _iter_mc_ind_area_files(file_output, freq): - data = df.data - # TODO: here we need to reshape the DF to comply with the schema (add index, add missing cols, - # reorder cols) + data = _adapt_df(variable_cols, df) writer.add_table(data.to_arrow()) diff --git a/tests/output/storage/v2/test_variable_storage.py b/tests/output/storage/v2/test_variable_storage.py new file mode 100644 index 0000000000..05726543d1 --- /dev/null +++ b/tests/output/storage/v2/test_variable_storage.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026, 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 polars as pl +from polars.testing import assert_frame_equal + +from antarest.output.filestudy.model import OutputDataFrame, VariableDescription +from antarest.output.storage.v2.variables_storage import _adapt_df + + +def test_adapt_df_to_columns() -> None: + columns = [ + VariableDescription("var1", None, None), + VariableDescription("var2", None, None), + VariableDescription("var3", None, None), + ] + output_df = OutputDataFrame( + data=pl.DataFrame( + [ + pl.Series(name="1", values=[0, 1], dtype=pl.Float64()), + pl.Series(name="2", values=[2, 3], dtype=pl.Float64()), + ] + ), + headers=[VariableDescription("var3", None, None), VariableDescription("var1", None, None)], + ) + adapted_df = _adapt_df(columns, output_df) + + # we expect to have second input column in 1st position (var1), then a null column for var2, + # then 1st input column in 3rd position for var3 + expected_df = pl.DataFrame( + [ + pl.Series(name="0", values=[2, 3], dtype=pl.Float64()), + pl.Series(name="1", values=[None, None], dtype=pl.Float64()), + pl.Series(name="2", values=[0, 1], dtype=pl.Float64()), + ] + ) + + assert_frame_equal(adapted_df, expected_df) From b2d121c580991a49792202d925f5173c0049c9a9 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 11:23:25 +0200 Subject: [PATCH 12/29] add support for index columns in writer Signed-off-by: Sylvain Leclerc --- antarest/core/serde/parquet_writer.py | 10 +- .../output/storage/v2/variables_storage.py | 160 ++++++++++++++---- tests/core/serde/test_parquet_writer.py | 8 +- .../storage/v2/test_variable_storage.py | 32 ++-- 4 files changed, 156 insertions(+), 54 deletions(-) diff --git a/antarest/core/serde/parquet_writer.py b/antarest/core/serde/parquet_writer.py index d4a5e3d5f7..f3bbf46aa3 100644 --- a/antarest/core/serde/parquet_writer.py +++ b/antarest/core/serde/parquet_writer.py @@ -55,7 +55,7 @@ def __enter__(self) -> "BatchParquetWriter": def __exit__(self, *args: Any, **kwargs: Any) -> None: self.close() - def add_table(self, table: pa.Table) -> None: + def append_table(self, table: pa.Table) -> None: if self._closed: raise ValueError("Writer is closed") self._current_batch.append(table) @@ -115,7 +115,7 @@ def write_dataframes_in_parquet_format_by_column_sets( first_df = _adapt_polars_schema(first_df) table = first_df.to_arrow() current_writer = BatchParquetWriter(file_path, table.schema) - current_writer.add_table(table) + current_writer.append_table(table) while True: try: @@ -143,7 +143,7 @@ def write_dataframes_in_parquet_format_by_column_sets( current_writer = BatchParquetWriter(file_path, table.schema) - current_writer.add_table(table) + current_writer.append_table(table) except StopIteration: return file_paths, new_index @@ -181,7 +181,7 @@ def write_dataframes_stream_parquet(path: Path, dataframes: Iterator[pd.DataFram raise ValueError("No dataframe provided") with BatchParquetWriter(path, schema) as writer: - writer.add_table(first_table) + writer.append_table(first_table) for df in dataframes: table = pa.Table.from_pandas(df) - writer.add_table(table) + writer.append_table(table) diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 2d0f004a68..13e5347d41 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -22,14 +22,14 @@ import shutil import tempfile from collections.abc import Iterator, Sequence +from dataclasses import dataclass from pathlib import Path -from typing import Any, Iterable +from typing import Any, Literal, Self, TypeAlias import polars as pl import polars.selectors as pls import pyarrow as pa from polars import Float64 -from pyarrow.lib import Field from antarest.core.exceptions import MCRootNotHandled, OutputAggregationError, OutputNotFound, OutputSubFolderNotFound from antarest.core.serde.parquet_writer import ( @@ -39,6 +39,7 @@ yield_dataframes_from_parquet, ) from antarest.output.filestudy.aggregation import AggregatorManager +from antarest.output.filestudy.iteration import OutputFileData, iterate_output_data from antarest.output.filestudy.matrixfiles import get_start_column, parse_output_file from antarest.output.filestudy.model import ( MCYEAR_COL, @@ -49,7 +50,6 @@ MCIndAreasQueryFile, MCIndLinksQueryFile, MCRoot, - OutputDataFrame, QueryFileType, VariableDescription, find_mode_dir, @@ -135,34 +135,130 @@ def _aggregate_to_parquet( _merge_intermediate_parquets(file_paths, new_index, target_path) -def _iter_mc_ind_area_files( - file_output: FileOutput, freq: MatrixFrequency -) -> Iterable[OutputDataFrame[VariableDescription]]: - start_col = get_start_column(freq) - for mc_year in file_output.mc_years: - for area_id in file_output.mc_ind_area_ids: - if data_file := file_output.get_mc_ind_file(mc_year, MCIndAreasQueryFile.VALUES, area_id, freq): - yield parse_output_file(data_file, start_col) +IndexCol: TypeAlias = Literal["mcYear", "area", "timeId"] +# Mapping to pyarrow fields +INDEX_FIELDS: dict[IndexCol, pa.Field[Any]] = { + "mcYear": pa.field("mcYear", pa.int32()), + "area": pa.field("area", pa.large_string()), # polars uses large_string and not just string + "timeId": pa.field("timeId", pa.int32()), +} -def _adapt_df(variable_cols: list[VariableDescription], df: OutputDataFrame[VariableDescription]) -> pl.DataFrame: + +@dataclass(frozen=True) +class IndexedOutputDataFrame: + """ + An output dataframe with columns for variables AND for "index" data such as the timestep, the element ID, + the MC year. """ - Reshapes the DF so that the column order matches the list of variables, and fills missing variables with null cols - TODO: create a writer class which encapsulates this plus the writing ? + index_cols: Sequence[IndexCol] + var_cols: Sequence[VariableDescription] + + data: pl.DataFrame + + +class ParquetOutputWriter: """ - col_for_variable = {v: i for i, v in enumerate(df.headers)} - nulls = pl.lit(None, dtype=Float64()) - res = df.data.select( - [ - pls.by_index(col_for_variable[v]).alias(str(i)) if v in col_for_variable else nulls.alias(str(i)) - for i, v in enumerate(variable_cols) + Utility class to append polars DF to a parquet file, taking care of adapting it to the required schema + (list of columns), and grouping them in not too small row groups (through batch parquet writer). + + Columns are named after variables but mainly for debugging purpose: the source of truth for variable columns + metadata remains the information stored in database. + """ + + def __init__(self, target_path: Path, index_cols: list[IndexCol], var_cols: Sequence[VariableDescription]) -> None: + self.index_cols = index_cols + self.var_cols = var_cols + self.writer = BatchParquetWriter(target_path, schema=self._create_schema()) + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: Any, **kwargs: Any) -> None: + self.writer.close() + + def _create_schema(self) -> pa.Schema: + return pa.schema( + [INDEX_FIELDS[c] for c in self.index_cols] + + [pa.field(self._col_name(i), pa.float64()) for i in range(len(self.var_cols))] + ) + + def _col_name(self, index: int) -> str: + """ + Trying to have a meaningful naming mainly for debugging purpose. + For business logic, the code MUST rely on database metadata instead. + """ + var = self.var_cols[index] + return f"{var.name} {var.statistic_type}" if var.statistic_type else var.name + + def _adapt_df(self, output_df: IndexedOutputDataFrame) -> pa.Table: + offset = len(self.index_cols) + col_for_variable = {v: offset + i for i, v in enumerate(output_df.var_cols)} + nulls = pl.lit(None, dtype=Float64()) + # just keep the index cols as is (just enforcing the "right" name here + index_cols = [pls.by_index(i).alias(name) for i, name in enumerate(self.index_cols)] + # reorder var cols and insert nulls where missing + var_cols = [ + pls.by_index(col_for_variable[v]).alias(self._col_name(i)) + if v in col_for_variable + else nulls.alias(self._col_name(i)) + for i, v in enumerate(self.var_cols) ] - ) - return res + adapted = output_df.data.select(index_cols + var_cols) + return adapted.to_arrow() + def append_output_df(self, output_df: IndexedOutputDataFrame) -> None: + if output_df.index_cols != self.index_cols: + raise ValueError( + f"Dataframe index differs from parquet file index ({output_df.index_cols} != {self.index_cols})" + ) + self.writer.append_table(self._adapt_df(output_df)) -def _extract_areas( + +_TIME_COL = pl.int_range(pl.len(), dtype=pl.Int32()) + + +def time_col() -> pl.Expr: + return _TIME_COL + + +def element_id_col(element_id: str) -> pl.Expr: + return pl.lit(element_id, dtype=pl.String()) + + +def mc_year_col(mc_year: int | Literal["mc-all"]) -> pl.Expr: + if mc_year == "mc-all": + raise ValueError("Should not created time id col for mc-all dataframe") + return pl.lit(mc_year, dtype=pl.Int32()) + + +def index_df(data: OutputFileData) -> IndexedOutputDataFrame: + """ + Adds index columns (mc year, element identifier(s), ) to raw variables dataframes + """ + metadata = data.file.metadata + df = data.data + match metadata.file_type: + case MCIndAreasQueryFile.VALUES: + return IndexedOutputDataFrame( + index_cols=["mcYear", "area", "timeId"], + var_cols=df.headers, + data=df.data.select( + mc_year_col(metadata.year), element_id_col(metadata.element_id), time_col(), pl.all() + ), + ) + case MCAllAreasQueryFile.VALUES: + return IndexedOutputDataFrame( + index_cols=["area", "timeId"], + var_cols=df.headers, + data=df.data.select(element_id_col(metadata.element_id), time_col(), pl.all()), + ) + + raise NotImplementedError(f"Not yet implemented: {metadata.file_type}") + + +def _extract_areas_refacto( index: VariablesIndex, file_output: FileOutput, target_dir: Path, @@ -170,19 +266,15 @@ def _extract_areas( variable_cols = index.get_variables("mc-ind", "area") - index_fields: list[Field[Any]] = [ - pa.field("area", pa.string()), - pa.field("mcYear", pa.int32()), - pa.field("timeId", pa.int32()), - ] - schema = pa.schema(index_fields + [pa.field(str(i), pa.float64()) for i in range(len(variable_cols))]) - for freq in MatrixFrequency: output_file_path = target_dir / f"mc-ind_areas_{freq.value}.parquet" - with BatchParquetWriter(output_file_path, schema) as writer: - for df in _iter_mc_ind_area_files(file_output, freq): - data = _adapt_df(variable_cols, df) - writer.add_table(data.to_arrow()) + with ParquetOutputWriter( + output_file_path, index_cols=["area", "mcYear", "timeId"], var_cols=variable_cols + ) as writer: + file_data = iterate_output_data(file_output.output_dir, MCIndAreasQueryFile.VALUES, freq, [], []) + indexed_dfs = map(index_df, file_data) + for df in indexed_dfs: + writer.append_output_df(df) # TODO: see above, this implementation needs to be replaced with one that uses the diff --git a/tests/core/serde/test_parquet_writer.py b/tests/core/serde/test_parquet_writer.py index 008cfd17a3..83c0e24587 100644 --- a/tests/core/serde/test_parquet_writer.py +++ b/tests/core/serde/test_parquet_writer.py @@ -100,7 +100,7 @@ def test_batch_parquet_writer_writes_one_batch(tmp_path: Path) -> None: ] with BatchParquetWriter(parquet_file, schema=tables[0].schema, row_group_size=5) as writer: for t in tables: - writer.add_table(t) + writer.append_table(t) with ParquetFile(parquet_file) as pf: assert pf.num_row_groups == 1 @@ -117,7 +117,7 @@ def test_batch_parquet_writer_writes_multiple_batches_when_size_exceeds_threshol ] with BatchParquetWriter(parquet_file, schema=tables[0].schema, row_group_size=5) as writer: for t in tables: - writer.add_table(t) + writer.append_table(t) with ParquetFile(parquet_file) as pf: assert pf.num_row_groups == 2 @@ -128,7 +128,7 @@ def test_batch_parquet_writer_cannot_add_table_to_closed_writer(tmp_path: Path) parquet_file = tmp_path / "file.parquet" table = pl.DataFrame(data=[(1, 2), (3, 4)], schema=["A", "B"], orient="row").to_arrow() with BatchParquetWriter(parquet_file, schema=table.schema, row_group_size=5) as writer: - writer.add_table(table) + writer.append_table(table) with pytest.raises(ValueError): - writer.add_table(table) + writer.append_table(table) diff --git a/tests/output/storage/v2/test_variable_storage.py b/tests/output/storage/v2/test_variable_storage.py index 05726543d1..ae37d4998d 100644 --- a/tests/output/storage/v2/test_variable_storage.py +++ b/tests/output/storage/v2/test_variable_storage.py @@ -9,37 +9,47 @@ # SPDX-License-Identifier: MPL-2.0 # # This file is part of the Antares project. +from pathlib import Path + import polars as pl from polars.testing import assert_frame_equal -from antarest.output.filestudy.model import OutputDataFrame, VariableDescription -from antarest.output.storage.v2.variables_storage import _adapt_df +from antarest.output.filestudy.model import VariableDescription +from antarest.output.storage.v2.variables_storage import IndexedOutputDataFrame, ParquetOutputWriter -def test_adapt_df_to_columns() -> None: - columns = [ - VariableDescription("var1", None, None), +def test_parquet_writer_adapts_df_to_columns(tmp_path: Path) -> None: + var_cols = [ + VariableDescription("var1", None, "exp"), VariableDescription("var2", None, None), VariableDescription("var3", None, None), ] - output_df = OutputDataFrame( + output_df = IndexedOutputDataFrame( data=pl.DataFrame( [ + pl.Series(name="area", values=["fr", "fr"], dtype=pl.String()), + pl.Series(name="timeId", values=[1, 2], dtype=pl.Int32()), pl.Series(name="1", values=[0, 1], dtype=pl.Float64()), pl.Series(name="2", values=[2, 3], dtype=pl.Float64()), ] ), - headers=[VariableDescription("var3", None, None), VariableDescription("var1", None, None)], + index_cols=["area", "timeId"], + var_cols=[VariableDescription("var3", None, None), VariableDescription("var1", None, "exp")], ) - adapted_df = _adapt_df(columns, output_df) + with ParquetOutputWriter(tmp_path / "output.parquet", ["area", "timeId"], var_cols) as writer: + writer.append_output_df(output_df) + + adapted_df = pl.read_parquet(tmp_path / "output.parquet") # we expect to have second input column in 1st position (var1), then a null column for var2, # then 1st input column in 3rd position for var3 expected_df = pl.DataFrame( [ - pl.Series(name="0", values=[2, 3], dtype=pl.Float64()), - pl.Series(name="1", values=[None, None], dtype=pl.Float64()), - pl.Series(name="2", values=[0, 1], dtype=pl.Float64()), + pl.Series(name="area", values=["fr", "fr"], dtype=pl.String()), + pl.Series(name="timeId", values=[1, 2], dtype=pl.Int32()), + pl.Series(name="var1 exp", values=[2, 3], dtype=pl.Float64()), + pl.Series(name="var2", values=[None, None], dtype=pl.Float64()), + pl.Series(name="var3", values=[0, 1], dtype=pl.Float64()), ] ) From fbddea82cbf2b23fdfcefb923ab69d0ad945ee7f Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 11:25:59 +0200 Subject: [PATCH 13/29] fix mypy related issue Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/variables_storage.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 13e5347d41..143028ce03 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -24,7 +24,7 @@ from collections.abc import Iterator, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Any, Literal, Self, TypeAlias +from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias import polars as pl import polars.selectors as pls @@ -137,8 +137,13 @@ def _aggregate_to_parquet( IndexCol: TypeAlias = Literal["mcYear", "area", "timeId"] +if TYPE_CHECKING: + Field = pa.Field[Any] +else: + Field = pa.Field + # Mapping to pyarrow fields -INDEX_FIELDS: dict[IndexCol, pa.Field[Any]] = { +INDEX_FIELDS: dict[IndexCol, Field] = { "mcYear": pa.field("mcYear", pa.int32()), "area": pa.field("area", pa.large_string()), # polars uses large_string and not just string "timeId": pa.field("timeId", pa.int32()), From 611465e6a96c59c5c67e07451cde6b217e91b69a Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 13:48:12 +0200 Subject: [PATCH 14/29] first test for parquet file creation Signed-off-by: Sylvain Leclerc --- .../output/storage/v2/variables_fetching.py | 12 ++-- .../output/storage/v2/variables_storage.py | 24 +++---- .../storage/v2/test_variable_storage.py | 64 ++++++++++++++++++- 3 files changed, 81 insertions(+), 19 deletions(-) diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index 5567ee2bc6..f00804e278 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -76,15 +76,19 @@ def _to_var_desc(db_var: DbParquetVariable) -> VariableDescription: return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) +def get_variables_index(session: Session, output_id: int) -> VariablesIndex: + output_variables = session.execute( + select(DbParquetVariable).where(DbParquetVariable.output_id == output_id) + ).scalars() + return VariablesIndex(output_variables) + + def get_area_variables( session: Session, output_id: int, aggregation: ScenarioAggregation, area_id: str ) -> list[VariableDescription]: # All variables, should load fast ? - output_variables = session.execute( - select(DbParquetVariable).where(DbParquetVariable.output_id == output_id) - ).scalars() - variables_index = VariablesIndex(output_variables) + variables_index = get_variables_index(session, output_id) # Get area information area = session.execute(select(DbParquetArea).where(DbParquetArea.area_id == area_id)).scalar_one() diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 143028ce03..6adc813b23 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -195,17 +195,17 @@ def _col_name(self, index: int) -> str: For business logic, the code MUST rely on database metadata instead. """ var = self.var_cols[index] - return f"{var.name} {var.statistic_type}" if var.statistic_type else var.name + return "__".join((p for p in (var.name, var.unit, var.statistic_type) if p)) def _adapt_df(self, output_df: IndexedOutputDataFrame) -> pa.Table: offset = len(self.index_cols) col_for_variable = {v: offset + i for i, v in enumerate(output_df.var_cols)} - nulls = pl.lit(None, dtype=Float64()) + nulls = pl.lit(None, dtype=Float64) # just keep the index cols as is (just enforcing the "right" name here index_cols = [pls.by_index(i).alias(name) for i, name in enumerate(self.index_cols)] # reorder var cols and insert nulls where missing var_cols = [ - pls.by_index(col_for_variable[v]).alias(self._col_name(i)) + pls.by_index(col_for_variable[v]).cast(dtype=Float64).alias(self._col_name(i)) if v in col_for_variable else nulls.alias(self._col_name(i)) for i, v in enumerate(self.var_cols) @@ -221,26 +221,26 @@ def append_output_df(self, output_df: IndexedOutputDataFrame) -> None: self.writer.append_table(self._adapt_df(output_df)) -_TIME_COL = pl.int_range(pl.len(), dtype=pl.Int32()) +_TIME_COL = pl.int_range(pl.len(), dtype=pl.Int32()).alias("timeId") def time_col() -> pl.Expr: return _TIME_COL -def element_id_col(element_id: str) -> pl.Expr: - return pl.lit(element_id, dtype=pl.String()) +def element_id_col(colname: str, element_id: str) -> pl.Expr: + return pl.lit(element_id, dtype=pl.String()).alias(colname) def mc_year_col(mc_year: int | Literal["mc-all"]) -> pl.Expr: if mc_year == "mc-all": raise ValueError("Should not created time id col for mc-all dataframe") - return pl.lit(mc_year, dtype=pl.Int32()) + return pl.lit(mc_year, dtype=pl.Int32()).alias("mcYear") def index_df(data: OutputFileData) -> IndexedOutputDataFrame: """ - Adds index columns (mc year, element identifier(s), ) to raw variables dataframes + Adds index columns (mc year, element identifier(s), ) to dataframes containing only variables values """ metadata = data.file.metadata df = data.data @@ -250,20 +250,20 @@ def index_df(data: OutputFileData) -> IndexedOutputDataFrame: index_cols=["mcYear", "area", "timeId"], var_cols=df.headers, data=df.data.select( - mc_year_col(metadata.year), element_id_col(metadata.element_id), time_col(), pl.all() + mc_year_col(metadata.year), element_id_col("area", metadata.element_id), time_col(), pl.all() ), ) case MCAllAreasQueryFile.VALUES: return IndexedOutputDataFrame( index_cols=["area", "timeId"], var_cols=df.headers, - data=df.data.select(element_id_col(metadata.element_id), time_col(), pl.all()), + data=df.data.select(element_id_col("area", metadata.element_id), time_col(), pl.all()), ) raise NotImplementedError(f"Not yet implemented: {metadata.file_type}") -def _extract_areas_refacto( +def extract_areas_refacto( index: VariablesIndex, file_output: FileOutput, target_dir: Path, @@ -274,7 +274,7 @@ def _extract_areas_refacto( for freq in MatrixFrequency: output_file_path = target_dir / f"mc-ind_areas_{freq.value}.parquet" with ParquetOutputWriter( - output_file_path, index_cols=["area", "mcYear", "timeId"], var_cols=variable_cols + output_file_path, index_cols=["mcYear", "area", "timeId"], var_cols=variable_cols ) as writer: file_data = iterate_output_data(file_output.output_dir, MCIndAreasQueryFile.VALUES, freq, [], []) indexed_dfs = map(index_df, file_data) diff --git a/tests/output/storage/v2/test_variable_storage.py b/tests/output/storage/v2/test_variable_storage.py index ae37d4998d..284a0f1316 100644 --- a/tests/output/storage/v2/test_variable_storage.py +++ b/tests/output/storage/v2/test_variable_storage.py @@ -12,10 +12,24 @@ from pathlib import Path import polars as pl +import pytest from polars.testing import assert_frame_equal +from sqlalchemy.orm import Session -from antarest.output.filestudy.model import VariableDescription -from antarest.output.storage.v2.variables_storage import IndexedOutputDataFrame, ParquetOutputWriter +from antarest.output.filestudy.model import FileOutput, VariableDescription +from antarest.output.storage.v2.dbmodel import DbParquetOutput +from antarest.output.storage.v2.variables_fetching import get_variables_index +from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database +from antarest.output.storage.v2.variables_storage import ( + IndexedOutputDataFrame, + ParquetOutputWriter, + extract_areas_refacto, +) + + +@pytest.fixture +def output_dir(data_dir: Path) -> Path: + return data_dir / "20260810-1420eco-thermal_groups" def test_parquet_writer_adapts_df_to_columns(tmp_path: Path) -> None: @@ -47,10 +61,54 @@ def test_parquet_writer_adapts_df_to_columns(tmp_path: Path) -> None: [ pl.Series(name="area", values=["fr", "fr"], dtype=pl.String()), pl.Series(name="timeId", values=[1, 2], dtype=pl.Int32()), - pl.Series(name="var1 exp", values=[2, 3], dtype=pl.Float64()), + pl.Series(name="var1__exp", values=[2, 3], dtype=pl.Float64()), pl.Series(name="var2", values=[None, None], dtype=pl.Float64()), pl.Series(name="var3", values=[0, 1], dtype=pl.Float64()), ] ) assert_frame_equal(adapted_df, expected_df) + + +def test_area_parquet_file_creation(output_dir: Path, db_session: Session, tmp_path: Path) -> None: + # TODO: should probably not need to go to the DB to get variables index ... + # TODO: should not get any file for empty frequencies + + db_output = DbParquetOutput(id=0) + db_session.add(db_output) + db_session.flush() + + file_output = FileOutput(output_dir) + extract_output_variables_to_database(db_session, db_output.id, file_output) + db_session.flush() + + extract_output_variables_to_database(db_session, 0, file_output) + index = get_variables_index(db_session, 0) + + target_dir = tmp_path / "output" + target_dir.mkdir() + + extract_areas_refacto(index, file_output, target_dir) + + monthly_file = target_dir / "mc-ind_areas_monthly.parquet" + assert monthly_file.is_file() + + df = pl.read_parquet(monthly_file) + + assert df.columns == [ + "mcYear", + "area", + "timeId", + "CO2 EMIS.__MWh", + "AVL DTG__MWh", + "DTG MRG__MWh", + "MAX MRG__MWh", + "NP COST__Euro", + "RES LOAD__MWh", + "NODU", + "ES_NUCLEAR_TH_PROD__MWh", + "FR_NUCLEAR_TH_PROD__MWh", + "CO2 EMIS.__Tons", + "ES_NUCLEAR__MWh", + "FR_NUCLEAR__MWh", + ] From 26b33951e98febd218d67428f7df3cbb7b15ead5 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 13:55:32 +0200 Subject: [PATCH 15/29] first tests for area parquet file creation Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/variables_storage.py | 8 ++++++-- tests/output/storage/v2/test_variable_storage.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 6adc813b23..1bdd04b31d 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -175,13 +175,15 @@ class ParquetOutputWriter: def __init__(self, target_path: Path, index_cols: list[IndexCol], var_cols: Sequence[VariableDescription]) -> None: self.index_cols = index_cols self.var_cols = var_cols - self.writer = BatchParquetWriter(target_path, schema=self._create_schema()) + self.target_path = target_path + self.writer: BatchParquetWriter | None = None def __enter__(self) -> Self: return self def __exit__(self, *args: Any, **kwargs: Any) -> None: - self.writer.close() + if self.writer: + self.writer.close() def _create_schema(self) -> pa.Schema: return pa.schema( @@ -218,6 +220,8 @@ def append_output_df(self, output_df: IndexedOutputDataFrame) -> None: raise ValueError( f"Dataframe index differs from parquet file index ({output_df.index_cols} != {self.index_cols})" ) + if not self.writer: + self.writer = BatchParquetWriter(self.target_path, schema=self._create_schema()) self.writer.append_table(self._adapt_df(output_df)) diff --git a/tests/output/storage/v2/test_variable_storage.py b/tests/output/storage/v2/test_variable_storage.py index 284a0f1316..b7e0d1d216 100644 --- a/tests/output/storage/v2/test_variable_storage.py +++ b/tests/output/storage/v2/test_variable_storage.py @@ -72,7 +72,6 @@ def test_parquet_writer_adapts_df_to_columns(tmp_path: Path) -> None: def test_area_parquet_file_creation(output_dir: Path, db_session: Session, tmp_path: Path) -> None: # TODO: should probably not need to go to the DB to get variables index ... - # TODO: should not get any file for empty frequencies db_output = DbParquetOutput(id=0) db_session.add(db_output) @@ -90,6 +89,7 @@ def test_area_parquet_file_creation(output_dir: Path, db_session: Session, tmp_p extract_areas_refacto(index, file_output, target_dir) + assert len(list(target_dir.iterdir())) == 1 monthly_file = target_dir / "mc-ind_areas_monthly.parquet" assert monthly_file.is_file() From 2ca742ca4f597c4c887dc82a17d468560604ff43 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 16:30:45 +0200 Subject: [PATCH 16/29] wip Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/dbmodel.py | 51 ++++---- antarest/output/storage/v2/download.py | 115 ++++++++++++++++++ .../output/storage/v2/variables_fetching.py | 7 ++ 3 files changed, 148 insertions(+), 25 deletions(-) create mode 100644 antarest/output/storage/v2/download.py diff --git a/antarest/output/storage/v2/dbmodel.py b/antarest/output/storage/v2/dbmodel.py index a1da0e94b0..87746eccfc 100644 --- a/antarest/output/storage/v2/dbmodel.py +++ b/antarest/output/storage/v2/dbmodel.py @@ -29,6 +29,29 @@ ScenarioAggregation: TypeAlias = Literal["mc-ind", "mc-all"] +class IntList(types.TypeDecorator[list[int]]): + """ + Stores a list of integers as a comma separated string. + + Can avoid many to many relationships which would not be useful. + """ + + impl = String + cache_ok = True + + @override + def process_bind_param(self, value: list[int] | None, dialect: Dialect) -> str: + if not isinstance(value, list): + raise ValueError("Expected a list of int") + return ",".join(str(c) for c in value) + + @override + def process_result_value(self, value: Any | None, dialect: Dialect) -> list[int]: + if not isinstance(value, str): + raise ValueError("Expected a string.") + return [int(c) for c in value.split(",")] + + class DbParquetOutput(Base): # TODO: we should merge the existing v2_output_metadata tables into this one # the integer identifier will be easier and more efficient to use than the couple of strings @@ -37,6 +60,7 @@ class DbParquetOutput(Base): __tablename__ = "parquet_output" id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True) + playlist: Mapped[list[int]] = mapped_column(IntList) class DbParquetVariable(Base): @@ -63,29 +87,6 @@ class DbParquetVariable(Base): statistic_type: Mapped[str | None] -class Columns(types.TypeDecorator[list[int]]): - """ - Stores a list of columns as a comma separated string. - - Avoids a many to many relationship which would not be useful. - """ - - impl = String - cache_ok = True - - @override - def process_bind_param(self, value: list[int] | None, dialect: Dialect) -> str: - if not isinstance(value, list): - raise ValueError("Expected a list of int for variable columns") - return ",".join(str(c) for c in value) - - @override - def process_result_value(self, value: Any | None, dialect: Dialect) -> list[int]: - if not isinstance(value, str): - raise ValueError("Expected a string in variable columns.") - return [int(c) for c in value.split(",")] - - class DbParquetArea(Base): """ Information related to an area of an output, in particular which variables it has data for, @@ -100,8 +101,8 @@ class DbParquetArea(Base): output_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) area_id: Mapped[str] = mapped_column(primary_key=True) - mc_all_vars: Mapped[list[int]] = mapped_column(Columns) - mc_ind_vars: Mapped[list[int]] = mapped_column(Columns) + mc_all_vars: Mapped[list[int]] = mapped_column(IntList) + mc_ind_vars: Mapped[list[int]] = mapped_column(IntList) # TODO: add tables for other element types: links, thermal clusters, etc diff --git a/antarest/output/storage/v2/download.py b/antarest/output/storage/v2/download.py new file mode 100644 index 0000000000..0a4245faa7 --- /dev/null +++ b/antarest/output/storage/v2/download.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026, 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. + + +""" +Support for the "download" API +""" + +import itertools +from pathlib import Path + +from polars import col, read_parquet, scan_parquet +from polars.selectors import by_index +from sqlalchemy import select +from sqlalchemy.orm import Session + +from antarest.output.filestudy.model import MCAllAreasQueryFile, MCIndAreasQueryFile, QueryFileType +from antarest.output.model import MatrixAggregationResultDTO, StudyDownloadDTO, StudyDownloadType, TimeSeriesData +from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetOutput, ElementType, ScenarioAggregation +from antarest.output.storage.v2.variables_fetching import get_variables_index +from antarest.study.model import MatrixFrequency + + +def _parquet_file_name(element_type: ElementType, frequency: MatrixFrequency) -> str: + match element_type: + case "area": + return f"mc-ind_areas_{frequency.value}.parquet" + case _: + raise NotImplementedError("Not yet implemented") + + +def _parquet_file(parquet_dir: Path, element_type: ElementType, frequency: MatrixFrequency) -> Path: + return parquet_dir / _parquet_file_name(element_type, frequency) + + +def build_matrix_aggregation_result( + session: Session, output_id: int, parquet_dir: Path, data_selection: StudyDownloadDTO +) -> MatrixAggregationResultDTO: + var_index = get_variables_index(session, output_id) + + element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system + + db_output = session.execute(select(DbParquetOutput).where(DbParquetOutput.id == output_id)).scalar_one() + mc_years = db_output.playlist + if data_selection.years: + mc_years = sorted(set(mc_years).intersection(data_selection.years)) + + if data_selection.type == StudyDownloadType.AREA: + parquet_file = _parquet_file(parquet_dir, "area", data_selection.level) + areas_df = scan_parquet(parquet_file) + offset = 3 # 3 index column. TODO: make it less fragile by storing it somewhere? or at least have a function + if data_selection.years: + areas_df = areas_df.filter(col("mcYear").is_in(data_selection.years)) + if data_selection.filter: + areas_df = areas_df.filter(col("area").is_in(data_selection.filter)) + if data_selection.columns: + area_vars = var_index.get_variables("mc-ind", "area") + selected_cols = [offset + c for c, v in enumerate(area_vars) if v.name in data_selection.columns] + areas_df = areas_df.select(by_index(selected_cols)) + + areas = session.execute(select(DbParquetArea).where(DbParquetArea.output_id == output_id)).scalars().fetchall() + + if data_selection.filter: + areas = [a for a in areas if a in data_selection.filter] + + areas = sorted(areas, key=lambda a: a.area_id) + + for a in itertools.product(areas, data_selection): + df = areas_df.filter(col("area") == a.area_id) + ts_data = element_results.setdefault( + a.area_id, TimeSeriesData(type=data_selection.type, name=a.area_id, data={}) + ) + numerical_data = df.to_series(var_index).cast(float).to_list() + ts_data.data.setdefault(str(year), []).append( + TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) + ) + + # Reshaping to target model + element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system + for file_data in output_data: + element_name = file_data.file.metadata.element_id + year = file_data.file.metadata.year + df = file_data.data.data + + # TODO: better handling of the link case, with 2 separate strings instead of that arbitrary formatting + if data_selection.type == StudyDownloadType.LINK: + element_name = "^".join(element_name.split(" - ")) + + for var_index, var in enumerate(file_data.data.headers): + if data_selection.columns and var.name not in data_selection.columns: + continue + ts_data = element_results.setdefault( + element_name, TimeSeriesData(type=data_selection.type, name=element_name, data={}) + ) + numerical_data = df.to_series(var_index).cast(float).to_list() + ts_data.data.setdefault(str(year), []).append( + TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) + ) + + time_index = get_start_date(None, output_path, data_selection.level) + return MatrixAggregationResultDTO( + index=time_index, + data=list(element_results.values()), + ) + + return MatrixAggregationResultDTO() diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index f00804e278..f22907b394 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from typing import Iterable, Sequence +from polars import DataFrame from sqlalchemy import select from sqlalchemy.orm import Session @@ -37,6 +38,12 @@ class VariableColumn: statistic_type: str | None +@dataclass(frozen=True) +class AreaVariables: + area_id: str + variables: Sequence[VariableColumn] + + class VariablesIndex: """ Helper class to retrieve variable info from DB models. From f72dd4c13c33abedf6ceb5a97a57f09e9540aae8 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 20:35:49 +0200 Subject: [PATCH 17/29] reorganize metadata implementation Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/download.py | 76 ++++------ antarest/output/storage/v2/metadata.py | 136 +++++++++++++++++ .../output/storage/v2/variables_fetching.py | 1 - tests/output/storage/v2/test_download.py | 139 ++++++++++++++++++ 4 files changed, 304 insertions(+), 48 deletions(-) create mode 100644 antarest/output/storage/v2/metadata.py create mode 100644 tests/output/storage/v2/test_download.py diff --git a/antarest/output/storage/v2/download.py b/antarest/output/storage/v2/download.py index 0a4245faa7..8c17c581e3 100644 --- a/antarest/output/storage/v2/download.py +++ b/antarest/output/storage/v2/download.py @@ -18,15 +18,20 @@ import itertools from pathlib import Path -from polars import col, read_parquet, scan_parquet +from polars import col, scan_parquet from polars.selectors import by_index -from sqlalchemy import select -from sqlalchemy.orm import Session -from antarest.output.filestudy.model import MCAllAreasQueryFile, MCIndAreasQueryFile, QueryFileType -from antarest.output.model import MatrixAggregationResultDTO, StudyDownloadDTO, StudyDownloadType, TimeSeriesData -from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetOutput, ElementType, ScenarioAggregation -from antarest.output.storage.v2.variables_fetching import get_variables_index +from antarest.output.model import ( + MatrixAggregationResultDTO, + StudyDownloadDTO, + StudyDownloadType, + TimeSerie, + TimeSeriesData, +) +from antarest.output.storage.v2.dbmodel import ( + ElementType, +) +from antarest.output.storage.v2.metadata import IParquetOutputMetadata from antarest.study.model import MatrixFrequency @@ -43,14 +48,15 @@ def _parquet_file(parquet_dir: Path, element_type: ElementType, frequency: Matri def build_matrix_aggregation_result( - session: Session, output_id: int, parquet_dir: Path, data_selection: StudyDownloadDTO + output_metadata: IParquetOutputMetadata, parquet_dir: Path, data_selection: StudyDownloadDTO ) -> MatrixAggregationResultDTO: - var_index = get_variables_index(session, output_id) + # TODO: works more or less but it's not quite clear. + + area_vars = output_metadata.get_variables("mc-ind", "area") element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system - db_output = session.execute(select(DbParquetOutput).where(DbParquetOutput.id == output_id)).scalar_one() - mc_years = db_output.playlist + mc_years = output_metadata.mc_years if data_selection.years: mc_years = sorted(set(mc_years).intersection(data_selection.years)) @@ -63,53 +69,29 @@ def build_matrix_aggregation_result( if data_selection.filter: areas_df = areas_df.filter(col("area").is_in(data_selection.filter)) if data_selection.columns: - area_vars = var_index.get_variables("mc-ind", "area") selected_cols = [offset + c for c, v in enumerate(area_vars) if v.name in data_selection.columns] areas_df = areas_df.select(by_index(selected_cols)) - areas = session.execute(select(DbParquetArea).where(DbParquetArea.output_id == output_id)).scalars().fetchall() + areas = output_metadata.mc_ind_areas if data_selection.filter: - areas = [a for a in areas if a in data_selection.filter] + areas = [a for a in areas if a.area_id in data_selection.filter] areas = sorted(areas, key=lambda a: a.area_id) - for a in itertools.product(areas, data_selection): - df = areas_df.filter(col("area") == a.area_id) - ts_data = element_results.setdefault( - a.area_id, TimeSeriesData(type=data_selection.type, name=a.area_id, data={}) - ) - numerical_data = df.to_series(var_index).cast(float).to_list() - ts_data.data.setdefault(str(year), []).append( - TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) - ) - - # Reshaping to target model - element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system - for file_data in output_data: - element_name = file_data.file.metadata.element_id - year = file_data.file.metadata.year - df = file_data.data.data - - # TODO: better handling of the link case, with 2 separate strings instead of that arbitrary formatting - if data_selection.type == StudyDownloadType.LINK: - element_name = "^".join(element_name.split(" - ")) - - for var_index, var in enumerate(file_data.data.headers): - if data_selection.columns and var.name not in data_selection.columns: - continue + for year, area in itertools.product(mc_years, areas): + df = areas_df.filter(col("area") == area.area_id).filter(col("mcYear") == year).collect() ts_data = element_results.setdefault( - element_name, TimeSeriesData(type=data_selection.type, name=element_name, data={}) - ) - numerical_data = df.to_series(var_index).cast(float).to_list() - ts_data.data.setdefault(str(year), []).append( - TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) + area.area_id, TimeSeriesData(type=data_selection.type, name=area.area_id, data={}) ) + for var_index in area.variables: + var = area_vars[var_index] + numerical_data = df.to_series(offset + var_index).cast(float).to_list() + ts_data.data.setdefault(str(year), []).append( + TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) + ) - time_index = get_start_date(None, output_path, data_selection.level) return MatrixAggregationResultDTO( - index=time_index, + index=output_metadata.get_time_index(data_selection.level), data=list(element_results.values()), ) - - return MatrixAggregationResultDTO() diff --git a/antarest/output/storage/v2/metadata.py b/antarest/output/storage/v2/metadata.py new file mode 100644 index 0000000000..5848e6d10e --- /dev/null +++ b/antarest/output/storage/v2/metadata.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026, 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 abc import ABC, abstractmethod +from dataclasses import dataclass +from functools import cached_property +from typing import Sequence + +from sqlalchemy import select +from sqlalchemy.orm import Session +from typing_extensions import override + +from antarest.output.filestudy.model import VariableDescription +from antarest.output.model import MatrixIndex +from antarest.output.storage.v2.dbmodel import ( + DbParquetArea, + DbParquetOutput, + DbParquetVariable, + ElementType, + ScenarioAggregation, +) +from antarest.output.storage.v2.variables_fetching import VariablesIndex +from antarest.study.model import MatrixFrequency + + +@dataclass(frozen=True) +class AreaVariables: + """ + Attributes: + area_id: Area identifier + variables: Indices of variables for that area, in the parquet file (does not include the index offset for now ...). + """ + + area_id: str + variables: Sequence[int] + + +class IParquetOutputMetadata(ABC): + """ + Centralizes access to all kind of metadata for one parquet output: + MC years, areas, variables ... + """ + + @property + @abstractmethod + def mc_years(self) -> list[int]: + """ + The list of MC years for which the output contains some data + """ + + @abstractmethod + def get_time_index(self, frequency: MatrixFrequency) -> MatrixIndex: + """ + Time index for the specified frequency + """ + + @abstractmethod + def get_variables( + self, aggregation: ScenarioAggregation, element_type: ElementType + ) -> Sequence[VariableDescription]: + """ + The list of variables for the specified mc-ind/mc-year aggregation and element_type, + in the same order as the corresponding columns in parquet files. + """ + + @property + @abstractmethod + def mc_ind_areas(self) -> Sequence[AreaVariables]: + """ + The list of areas for mc-ind results, and the corresponding variables for which they have data. + """ + + +class ParquetOuputMetadataImpl(IParquetOutputMetadata): + """ + Implementation which gets metadata from the DB, as lazily as possible. + + Caches retrieved data for re-use. + """ + + def __init__(self, session: Session, output_id: int) -> None: + self.session = session + self.output_id = output_id + + @cached_property + def db_output(self) -> DbParquetOutput: + return self.session.execute(select(DbParquetOutput).where(DbParquetOutput.id == self.output_id)).scalar_one() + + @cached_property + def db_areas(self) -> Sequence[DbParquetArea]: + return ( + self.session.execute(select(DbParquetArea).where(DbParquetArea.output_id == self.output_id)) + .scalars() + .fetchall() + ) + + @cached_property + def db_vars(self) -> Sequence[DbParquetVariable]: + return ( + self.session.execute(select(DbParquetVariable).where(DbParquetVariable.output_id == self.output_id)) + .scalars() + .fetchall() + ) + + @cached_property + def variables_index(self) -> VariablesIndex: + return VariablesIndex(self.db_vars) + + @override + def get_variables( + self, aggregation: ScenarioAggregation, element_type: ElementType + ) -> Sequence[VariableDescription]: + return self.variables_index.get_variables(aggregation, element_type) + + @override + @property + def mc_years(self) -> list[int]: + return self.db_output.playlist + + @override + def get_time_index(self, frequency: MatrixFrequency) -> MatrixIndex: + # TODO: get it from current implementation + return MatrixIndex() + + @override + @property + def mc_ind_areas(self) -> Sequence[AreaVariables]: + return tuple(AreaVariables(area_id=a.area_id, variables=a.mc_ind_vars) for a in self.db_areas) diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index f22907b394..91dfc33eef 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -17,7 +17,6 @@ from dataclasses import dataclass from typing import Iterable, Sequence -from polars import DataFrame from sqlalchemy import select from sqlalchemy.orm import Session diff --git a/tests/output/storage/v2/test_download.py b/tests/output/storage/v2/test_download.py new file mode 100644 index 0000000000..e2de59c113 --- /dev/null +++ b/tests/output/storage/v2/test_download.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026, 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 pathlib import Path + +import pytest +from sqlalchemy.orm import Session + +from antarest.output.filestudy.model import FileOutput +from antarest.output.model import StudyDownloadDTO, StudyDownloadType +from antarest.output.storage.v2.dbmodel import DbParquetOutput +from antarest.output.storage.v2.download import ( + build_matrix_aggregation_result, +) +from antarest.output.storage.v2.metadata import ParquetOuputMetadataImpl +from antarest.output.storage.v2.variables_fetching import get_variables_index +from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database +from antarest.output.storage.v2.variables_storage import extract_areas_refacto +from antarest.study.model import MatrixFrequency + + +@pytest.fixture +def output_dir(data_dir: Path) -> Path: + return data_dir / "20260810-1420eco-thermal_groups" + + +def test_download_areas(output_dir: Path, db_session: Session, tmp_path: Path) -> None: + # TODO: simplify setup + + # Setup + + db_output = DbParquetOutput(id=0, playlist=[1, 2]) + db_session.add(db_output) + db_session.flush() + + file_output = FileOutput(output_dir) + extract_output_variables_to_database(db_session, db_output.id, file_output) + db_session.flush() + + extract_output_variables_to_database(db_session, 0, file_output) + index = get_variables_index(db_session, 0) + + target_dir = tmp_path / "output" + target_dir.mkdir() + + extract_areas_refacto(index, file_output, target_dir) + + assert len(list(target_dir.iterdir())) == 1 + monthly_file = target_dir / "mc-ind_areas_monthly.parquet" + assert monthly_file.is_file() + + # Actual test + + output_metadata = ParquetOuputMetadataImpl(db_session, db_output.id) + data_selection = StudyDownloadDTO(type=StudyDownloadType.AREA, years=[], level=MatrixFrequency.MONTHLY, filter=[]) + aggregate = build_matrix_aggregation_result(output_metadata, target_dir, data_selection) + + year1_st_by_area = {data.name: data.data["1"] for data in aggregate.data} + es_variables = [ts.name for ts in year1_st_by_area["es"]] + fr_variables = [ts.name for ts in year1_st_by_area["fr"]] + + # Same test as for filestudy + assert es_variables == [ + "CO2 EMIS.", + "ES_NUCLEAR", # We only get ES group + "AVL DTG", + "DTG MRG", + "MAX MRG", + "NP COST", + "NODU", + "RES LOAD", + ] + assert fr_variables == [ + "CO2 EMIS.", + "FR_NUCLEAR", # we only get FR group + "AVL DTG", + "DTG MRG", + "MAX MRG", + "NP COST", + "NODU", + "RES LOAD", + ] + + +def test_download_district(output_dir: Path, db_session: Session, tmp_path: Path) -> None: + # TODO: simplify setup + # TODO: make it pass + + # Setup + + db_output = DbParquetOutput(id=0, playlist=[1, 2]) + db_session.add(db_output) + db_session.flush() + + file_output = FileOutput(output_dir) + extract_output_variables_to_database(db_session, db_output.id, file_output) + db_session.flush() + + extract_output_variables_to_database(db_session, 0, file_output) + index = get_variables_index(db_session, 0) + + target_dir = tmp_path / "output" + target_dir.mkdir() + + extract_areas_refacto(index, file_output, target_dir) + + assert len(list(target_dir.iterdir())) == 1 + monthly_file = target_dir / "mc-ind_areas_monthly.parquet" + assert monthly_file.is_file() + + # Actual test + + output_metadata = ParquetOuputMetadataImpl(db_session, db_output.id) + + data_selection = StudyDownloadDTO(type=StudyDownloadType.DISTRICT, years=[1], level=MatrixFrequency.MONTHLY) + aggregate = build_matrix_aggregation_result(output_metadata, target_dir, data_selection) + + year1_st_by_area = {data.name: data.data["1"] for data in aggregate.data} + all_areas_variables = [ts.name for ts in year1_st_by_area["@ all areas"]] + + assert all_areas_variables == [ + "CO2 EMIS.", + "AVL DTG", + "DTG MRG", + "MAX MRG", + "NP COST", + "RES LOAD", + "NODU", + "ES_NUCLEAR_TH_PROD", + "FR_NUCLEAR_TH_PROD", + ] From c12e2e4bb29ea169b0aaf74891bdbfa326df5f79 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Mon, 24 Aug 2026 20:37:25 +0200 Subject: [PATCH 18/29] renaming Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/dbmodel.py | 2 +- antarest/output/storage/v2/metadata.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/antarest/output/storage/v2/dbmodel.py b/antarest/output/storage/v2/dbmodel.py index 87746eccfc..a93680daa6 100644 --- a/antarest/output/storage/v2/dbmodel.py +++ b/antarest/output/storage/v2/dbmodel.py @@ -60,7 +60,7 @@ class DbParquetOutput(Base): __tablename__ = "parquet_output" id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True) - playlist: Mapped[list[int]] = mapped_column(IntList) + mc_years: Mapped[list[int]] = mapped_column(IntList) class DbParquetVariable(Base): diff --git a/antarest/output/storage/v2/metadata.py b/antarest/output/storage/v2/metadata.py index 5848e6d10e..9210a1a1fb 100644 --- a/antarest/output/storage/v2/metadata.py +++ b/antarest/output/storage/v2/metadata.py @@ -123,7 +123,7 @@ def get_variables( @override @property def mc_years(self) -> list[int]: - return self.db_output.playlist + return self.db_output.mc_years @override def get_time_index(self, frequency: MatrixFrequency) -> MatrixIndex: From 0ef2eefb3a7be3998c4f354c14fa86dd5a284b66 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 08:52:57 +0200 Subject: [PATCH 19/29] moving stuff Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/metadata.py | 3 + .../output/storage/v2/variables_storage.py | 316 +++++++++--------- 2 files changed, 165 insertions(+), 154 deletions(-) diff --git a/antarest/output/storage/v2/metadata.py b/antarest/output/storage/v2/metadata.py index 9210a1a1fb..4e7cb7602b 100644 --- a/antarest/output/storage/v2/metadata.py +++ b/antarest/output/storage/v2/metadata.py @@ -47,6 +47,9 @@ class IParquetOutputMetadata(ABC): """ Centralizes access to all kind of metadata for one parquet output: MC years, areas, variables ... + + Probably interesting to keep it as an interface for now, because if we want to create it in + a separate process at computation time, implementation will not rely on the database. """ @property diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 1bdd04b31d..733c273e8b 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -55,7 +55,7 @@ find_mode_dir, get_output_object_type, ) -from antarest.output.storage.v2.variables_fetching import VariablesIndex +from antarest.output.storage.v2.metadata import IParquetOutputMetadata from antarest.study.model import MatrixFrequency logger = logging.getLogger(__name__) @@ -135,159 +135,6 @@ def _aggregate_to_parquet( _merge_intermediate_parquets(file_paths, new_index, target_path) -IndexCol: TypeAlias = Literal["mcYear", "area", "timeId"] - -if TYPE_CHECKING: - Field = pa.Field[Any] -else: - Field = pa.Field - -# Mapping to pyarrow fields -INDEX_FIELDS: dict[IndexCol, Field] = { - "mcYear": pa.field("mcYear", pa.int32()), - "area": pa.field("area", pa.large_string()), # polars uses large_string and not just string - "timeId": pa.field("timeId", pa.int32()), -} - - -@dataclass(frozen=True) -class IndexedOutputDataFrame: - """ - An output dataframe with columns for variables AND for "index" data such as the timestep, the element ID, - the MC year. - """ - - index_cols: Sequence[IndexCol] - var_cols: Sequence[VariableDescription] - - data: pl.DataFrame - - -class ParquetOutputWriter: - """ - Utility class to append polars DF to a parquet file, taking care of adapting it to the required schema - (list of columns), and grouping them in not too small row groups (through batch parquet writer). - - Columns are named after variables but mainly for debugging purpose: the source of truth for variable columns - metadata remains the information stored in database. - """ - - def __init__(self, target_path: Path, index_cols: list[IndexCol], var_cols: Sequence[VariableDescription]) -> None: - self.index_cols = index_cols - self.var_cols = var_cols - self.target_path = target_path - self.writer: BatchParquetWriter | None = None - - def __enter__(self) -> Self: - return self - - def __exit__(self, *args: Any, **kwargs: Any) -> None: - if self.writer: - self.writer.close() - - def _create_schema(self) -> pa.Schema: - return pa.schema( - [INDEX_FIELDS[c] for c in self.index_cols] - + [pa.field(self._col_name(i), pa.float64()) for i in range(len(self.var_cols))] - ) - - def _col_name(self, index: int) -> str: - """ - Trying to have a meaningful naming mainly for debugging purpose. - For business logic, the code MUST rely on database metadata instead. - """ - var = self.var_cols[index] - return "__".join((p for p in (var.name, var.unit, var.statistic_type) if p)) - - def _adapt_df(self, output_df: IndexedOutputDataFrame) -> pa.Table: - offset = len(self.index_cols) - col_for_variable = {v: offset + i for i, v in enumerate(output_df.var_cols)} - nulls = pl.lit(None, dtype=Float64) - # just keep the index cols as is (just enforcing the "right" name here - index_cols = [pls.by_index(i).alias(name) for i, name in enumerate(self.index_cols)] - # reorder var cols and insert nulls where missing - var_cols = [ - pls.by_index(col_for_variable[v]).cast(dtype=Float64).alias(self._col_name(i)) - if v in col_for_variable - else nulls.alias(self._col_name(i)) - for i, v in enumerate(self.var_cols) - ] - adapted = output_df.data.select(index_cols + var_cols) - return adapted.to_arrow() - - def append_output_df(self, output_df: IndexedOutputDataFrame) -> None: - if output_df.index_cols != self.index_cols: - raise ValueError( - f"Dataframe index differs from parquet file index ({output_df.index_cols} != {self.index_cols})" - ) - if not self.writer: - self.writer = BatchParquetWriter(self.target_path, schema=self._create_schema()) - self.writer.append_table(self._adapt_df(output_df)) - - -_TIME_COL = pl.int_range(pl.len(), dtype=pl.Int32()).alias("timeId") - - -def time_col() -> pl.Expr: - return _TIME_COL - - -def element_id_col(colname: str, element_id: str) -> pl.Expr: - return pl.lit(element_id, dtype=pl.String()).alias(colname) - - -def mc_year_col(mc_year: int | Literal["mc-all"]) -> pl.Expr: - if mc_year == "mc-all": - raise ValueError("Should not created time id col for mc-all dataframe") - return pl.lit(mc_year, dtype=pl.Int32()).alias("mcYear") - - -def index_df(data: OutputFileData) -> IndexedOutputDataFrame: - """ - Adds index columns (mc year, element identifier(s), ) to dataframes containing only variables values - """ - metadata = data.file.metadata - df = data.data - match metadata.file_type: - case MCIndAreasQueryFile.VALUES: - return IndexedOutputDataFrame( - index_cols=["mcYear", "area", "timeId"], - var_cols=df.headers, - data=df.data.select( - mc_year_col(metadata.year), element_id_col("area", metadata.element_id), time_col(), pl.all() - ), - ) - case MCAllAreasQueryFile.VALUES: - return IndexedOutputDataFrame( - index_cols=["area", "timeId"], - var_cols=df.headers, - data=df.data.select(element_id_col("area", metadata.element_id), time_col(), pl.all()), - ) - - raise NotImplementedError(f"Not yet implemented: {metadata.file_type}") - - -def extract_areas_refacto( - index: VariablesIndex, - file_output: FileOutput, - target_dir: Path, -) -> None: - - variable_cols = index.get_variables("mc-ind", "area") - - for freq in MatrixFrequency: - output_file_path = target_dir / f"mc-ind_areas_{freq.value}.parquet" - with ParquetOutputWriter( - output_file_path, index_cols=["mcYear", "area", "timeId"], var_cols=variable_cols - ) as writer: - file_data = iterate_output_data(file_output.output_dir, MCIndAreasQueryFile.VALUES, freq, [], []) - indexed_dfs = map(index_df, file_data) - for df in indexed_dfs: - writer.append_output_df(df) - - -# TODO: see above, this implementation needs to be replaced with one that uses the -# column indices that have been determined when parsing variable metadata def _extract_areas( output_dir: Path, base_path: Path, @@ -555,3 +402,164 @@ def read_output_from_parquet( if district_ids: parquet_path = target_dir / _parquet_file_name(mc_root, "districts", frequency) yield from _read_filtered(parquet_path, id_col, district_ids, mc_root, mc_years, columns_names, is_details) + + +# TODO: the implementation above needs to be replaced with one that uses the +# column indices that have been determined when parsing variable metadata + +IndexCol: TypeAlias = Literal["mcYear", "area", "timeId"] + +if TYPE_CHECKING: + Field = pa.Field[Any] +else: + Field = pa.Field + +# Mapping to pyarrow fields +INDEX_FIELDS: dict[IndexCol, Field] = { + "mcYear": pa.field("mcYear", pa.int32()), + "area": pa.field("area", pa.large_string()), # polars uses large_string and not just string + "timeId": pa.field("timeId", pa.int32()), +} + + +@dataclass(frozen=True) +class IndexedOutputDataFrame: + """ + An output dataframe with columns for variables AND for "index" data such as the timestep, the element ID, + the MC year. + """ + + index_cols: Sequence[IndexCol] + var_cols: Sequence[VariableDescription] + + data: pl.DataFrame + + +class ParquetOutputWriter: + """ + Utility class to append polars DF to a parquet file, taking care of adapting it to the required schema + (list of columns), and grouping them in not too small row groups (through batch parquet writer). + + Columns are named after variables but mainly for debugging purpose: the source of truth for variable columns + metadata remains the information stored in database. + """ + + def __init__(self, target_path: Path, index_cols: list[IndexCol], var_cols: Sequence[VariableDescription]) -> None: + self.index_cols = index_cols + self.var_cols = var_cols + self.target_path = target_path + self.writer: BatchParquetWriter | None = None + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: Any, **kwargs: Any) -> None: + if self.writer: + self.writer.close() + + def _create_schema(self) -> pa.Schema: + return pa.schema( + [INDEX_FIELDS[c] for c in self.index_cols] + + [pa.field(self._col_name(i), pa.float64()) for i in range(len(self.var_cols))] + ) + + def _col_name(self, index: int) -> str: + """ + Trying to have a meaningful naming mainly for debugging purpose. + For business logic, the code MUST rely on database metadata instead. + """ + var = self.var_cols[index] + return "__".join((p for p in (var.name, var.unit, var.statistic_type) if p)) + + def _adapt_df(self, output_df: IndexedOutputDataFrame) -> pa.Table: + offset = len(self.index_cols) + col_for_variable = {v: offset + i for i, v in enumerate(output_df.var_cols)} + nulls = pl.lit(None, dtype=Float64) + # just keep the index cols as is (just enforcing the "right" name here + index_cols = [pls.by_index(i).alias(name) for i, name in enumerate(self.index_cols)] + # reorder var cols and insert nulls where missing + var_cols = [ + pls.by_index(col_for_variable[v]).cast(dtype=Float64).alias(self._col_name(i)) + if v in col_for_variable + else nulls.alias(self._col_name(i)) + for i, v in enumerate(self.var_cols) + ] + adapted = output_df.data.select(index_cols + var_cols) + return adapted.to_arrow() + + def append_output_df(self, output_df: IndexedOutputDataFrame) -> None: + if output_df.index_cols != self.index_cols: + raise ValueError( + f"Dataframe index differs from parquet file index ({output_df.index_cols} != {self.index_cols})" + ) + if not self.writer: + self.writer = BatchParquetWriter(self.target_path, schema=self._create_schema()) + self.writer.append_table(self._adapt_df(output_df)) + + +_TIME_COL = pl.int_range(pl.len(), dtype=pl.Int32()).alias("timeId") + + +def time_col() -> pl.Expr: + return _TIME_COL + + +def element_id_col(colname: str, element_id: str) -> pl.Expr: + return pl.lit(element_id, dtype=pl.String()).alias(colname) + + +def mc_year_col(mc_year: int | Literal["mc-all"]) -> pl.Expr: + if mc_year == "mc-all": + raise ValueError("Should not created time id col for mc-all dataframe") + return pl.lit(mc_year, dtype=pl.Int32()).alias("mcYear") + + +def index_df(data: OutputFileData) -> IndexedOutputDataFrame: + """ + Adds index columns (mc year, element identifier(s), ) to dataframes containing only variables values + """ + metadata = data.file.metadata + df = data.data + match metadata.file_type: + case MCIndAreasQueryFile.VALUES: + return IndexedOutputDataFrame( + index_cols=["mcYear", "area", "timeId"], + var_cols=df.headers, + data=df.data.select( + mc_year_col(metadata.year), element_id_col("area", metadata.element_id), time_col(), pl.all() + ), + ) + case MCAllAreasQueryFile.VALUES: + return IndexedOutputDataFrame( + index_cols=["area", "timeId"], + var_cols=df.headers, + data=df.data.select(element_id_col("area", metadata.element_id), time_col(), pl.all()), + ) + + raise NotImplementedError(f"Not yet implemented: {metadata.file_type}") + + +def extract_areas_refacto( + metadata: IParquetOutputMetadata, + file_output: FileOutput, + target_dir: Path, +) -> None: + + variable_cols = metadata.get_variables("mc-ind", "area") + + for freq in MatrixFrequency: + output_file_path = target_dir / f"mc-ind_areas_{freq.value}.parquet" + with ParquetOutputWriter( + output_file_path, index_cols=["mcYear", "area", "timeId"], var_cols=variable_cols + ) as writer: + file_data = iterate_output_data(file_output.output_dir, MCIndAreasQueryFile.VALUES, freq, [], []) + indexed_dfs = map(index_df, file_data) + for df in indexed_dfs: + writer.append_output_df(df) + + +def create_parquet_files(metadata: IParquetOutputMetadata, file_output: FileOutput, target_dir: Path) -> None: + """ + Creates parquet files in target_dir in consistence with columns that have been defined in the metadata object. + """ + extract_areas_refacto(metadata, file_output, target_dir) From b9d1c14df8c4c1344748879836310305b8b3433b Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 09:03:23 +0200 Subject: [PATCH 20/29] clarify functions chaining Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/repository.py | 8 ++++++++ antarest/output/storage/v2/storage.py | 19 ++++++++++++++----- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/antarest/output/storage/v2/repository.py b/antarest/output/storage/v2/repository.py index b0dcccf4d7..cd79fd4d56 100644 --- a/antarest/output/storage/v2/repository.py +++ b/antarest/output/storage/v2/repository.py @@ -10,6 +10,7 @@ # # This file is part of the Antares project. from collections.abc import Iterator +from typing import Iterable from sqlalchemy import Boolean, Column, ForeignKeyConstraint, Integer, String, Table, delete, select from sqlalchemy.orm import Mapped, Session, mapped_column @@ -19,6 +20,7 @@ from antarest.core.utils.sql_utils import upsert_one from antarest.launcher.model import LogType from antarest.output.model import OutputVariablesList +from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable class DbOutputMetadataV2(Base): @@ -174,3 +176,9 @@ def save_output_variables_list(self, study_id: str, output_id: str, variables_li db_model = DbOutputVariablesV2.from_model(study_id, output_id, variables_list) self.session.add(db_model) self.session.commit() + + def save_variables(self, variables: Iterable[DbParquetVariable]) -> None: + self.session.add_all(variables) + + def save_areas(self, areas: Iterable[DbParquetArea]) -> None: + self.session.add_all(areas) diff --git a/antarest/output/storage/v2/storage.py b/antarest/output/storage/v2/storage.py index b2e8b098cc..8943671d7d 100644 --- a/antarest/output/storage/v2/storage.py +++ b/antarest/output/storage/v2/storage.py @@ -34,6 +34,7 @@ extract_archive_from_path, extract_archive_from_stream, ) +from antarest.core.utils.fastapi_sqlalchemy import db from antarest.core.utils.sqlalchemy import clone_orm_object from antarest.core.utils.utils import StopWatch from antarest.launcher.adapters.abstractlauncher import SimulationLogs @@ -43,7 +44,7 @@ from antarest.output.filestudy.metadata import ( extract_output_details, ) -from antarest.output.filestudy.model import QueryFileType +from antarest.output.filestudy.model import FileOutput, QueryFileType from antarest.output.filestudy.variables import extract_variables_list from antarest.output.model import MatrixAggregationResultDTO, OutputVariablesList, StudyDownloadDTO from antarest.output.model.download import MatrixIndex @@ -53,11 +54,14 @@ OutputMetadata, OutputStorageType, ) +from antarest.output.storage.v2.metadata import ParquetOuputMetadataImpl from antarest.output.storage.v2.repository import ( DbOutputMetadataV2, OutputV2Repository, ) +from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database from antarest.output.storage.v2.variables_storage import ( + create_parquet_files, extract_output_to_parquet, parquet_output_dir, read_output_from_parquet, @@ -222,10 +226,6 @@ def import_output( simulation_range = _extract_simulation_range(dir_path) - # TODO: first, extract variables metadata to database - variables_target = parquet_output_dir(self._variables_dir, study_id, output_name) - extract_output_to_parquet(dir_path, variables_target) - self._repository.save_output_metadata( DbOutputMetadataV2( study_id=study_id, @@ -244,6 +244,15 @@ def import_output( ) ) + # TODO: first, extract variables metadata to database + file_output = FileOutput(dir_path) + output_id = 0 # TODO: create it first with the metadata above + extract_output_variables_to_database(db.session, output_id, file_output) + + variables_target = parquet_output_dir(self._variables_dir, study_id, output_name) + metadata = ParquetOuputMetadataImpl(db.session, output_id) + create_parquet_files(metadata, file_output, variables_target) # TODO: complete implementation + self._save_logs(study_id, output_name, logs, dir_path) variables_list = extract_variables_list(dir_path) From ac9c9e324c2e7d3719a09b6ad3aa976fefce71df Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 09:07:09 +0200 Subject: [PATCH 21/29] add call site in storage implementation Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/storage.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/antarest/output/storage/v2/storage.py b/antarest/output/storage/v2/storage.py index 8943671d7d..f6b5f61854 100644 --- a/antarest/output/storage/v2/storage.py +++ b/antarest/output/storage/v2/storage.py @@ -54,6 +54,7 @@ OutputMetadata, OutputStorageType, ) +from antarest.output.storage.v2.download import build_matrix_aggregation_result from antarest.output.storage.v2.metadata import ParquetOuputMetadataImpl from antarest.output.storage.v2.repository import ( DbOutputMetadataV2, @@ -459,4 +460,8 @@ def get_original_file(self, study_id: str, output_id: str, url: list[str]) -> Or def get_matrix_aggregation_result( self, study_id: str, output_id: str, data_selection: StudyDownloadDTO ) -> MatrixAggregationResultDTO: - raise NotImplementedError() + metadata = self._require_metadata(study_id, output_id) + db_id = 0 # TODO: get from metadata + parquet_metadata = ParquetOuputMetadataImpl(db.session, db_id) + output_dir = parquet_output_dir(self._variables_dir, study_id, output_id) + return build_matrix_aggregation_result(parquet_metadata, output_dir, data_selection) From 9e028547f62b9f0381b2300a8fe990179b520a0a Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 09:08:09 +0200 Subject: [PATCH 22/29] remove comment Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/storage.py | 1 - 1 file changed, 1 deletion(-) diff --git a/antarest/output/storage/v2/storage.py b/antarest/output/storage/v2/storage.py index f6b5f61854..e331de3ae8 100644 --- a/antarest/output/storage/v2/storage.py +++ b/antarest/output/storage/v2/storage.py @@ -245,7 +245,6 @@ def import_output( ) ) - # TODO: first, extract variables metadata to database file_output = FileOutput(dir_path) output_id = 0 # TODO: create it first with the metadata above extract_output_variables_to_database(db.session, output_id, file_output) From 2fde291f387b9e9e5352484e3727cf7a136155dc Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 10:15:04 +0200 Subject: [PATCH 23/29] separate iteration logic and fix it Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/download.py | 64 ++++------------ antarest/output/storage/v2/iteration.py | 98 +++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 49 deletions(-) create mode 100644 antarest/output/storage/v2/iteration.py diff --git a/antarest/output/storage/v2/download.py b/antarest/output/storage/v2/download.py index 8c17c581e3..b8759e8708 100644 --- a/antarest/output/storage/v2/download.py +++ b/antarest/output/storage/v2/download.py @@ -15,12 +15,8 @@ Support for the "download" API """ -import itertools from pathlib import Path -from polars import col, scan_parquet -from polars.selectors import by_index - from antarest.output.model import ( MatrixAggregationResultDTO, StudyDownloadDTO, @@ -28,65 +24,35 @@ TimeSerie, TimeSeriesData, ) -from antarest.output.storage.v2.dbmodel import ( - ElementType, -) +from antarest.output.storage.v2.iteration import iterate_areas_df from antarest.output.storage.v2.metadata import IParquetOutputMetadata -from antarest.study.model import MatrixFrequency - - -def _parquet_file_name(element_type: ElementType, frequency: MatrixFrequency) -> str: - match element_type: - case "area": - return f"mc-ind_areas_{frequency.value}.parquet" - case _: - raise NotImplementedError("Not yet implemented") - - -def _parquet_file(parquet_dir: Path, element_type: ElementType, frequency: MatrixFrequency) -> Path: - return parquet_dir / _parquet_file_name(element_type, frequency) def build_matrix_aggregation_result( output_metadata: IParquetOutputMetadata, parquet_dir: Path, data_selection: StudyDownloadDTO ) -> MatrixAggregationResultDTO: - # TODO: works more or less but it's not quite clear. area_vars = output_metadata.get_variables("mc-ind", "area") element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system - mc_years = output_metadata.mc_years - if data_selection.years: - mc_years = sorted(set(mc_years).intersection(data_selection.years)) - - if data_selection.type == StudyDownloadType.AREA: - parquet_file = _parquet_file(parquet_dir, "area", data_selection.level) - areas_df = scan_parquet(parquet_file) - offset = 3 # 3 index column. TODO: make it less fragile by storing it somewhere? or at least have a function - if data_selection.years: - areas_df = areas_df.filter(col("mcYear").is_in(data_selection.years)) - if data_selection.filter: - areas_df = areas_df.filter(col("area").is_in(data_selection.filter)) - if data_selection.columns: - selected_cols = [offset + c for c, v in enumerate(area_vars) if v.name in data_selection.columns] - areas_df = areas_df.select(by_index(selected_cols)) - - areas = output_metadata.mc_ind_areas - - if data_selection.filter: - areas = [a for a in areas if a.area_id in data_selection.filter] - - areas = sorted(areas, key=lambda a: a.area_id) - - for year, area in itertools.product(mc_years, areas): - df = areas_df.filter(col("area") == area.area_id).filter(col("mcYear") == year).collect() + if data_selection.type in {StudyDownloadType.AREA, StudyDownloadType.DISTRICT}: + area_dfs = iterate_areas_df( + output_metadata, + parquet_dir, + data_selection.level, + data_selection.years, + data_selection.filter, + data_selection.columns, + ) + for area_df in area_dfs: + year, area_id, df = area_df.year, area_df.area_id, area_df.data ts_data = element_results.setdefault( - area.area_id, TimeSeriesData(type=data_selection.type, name=area.area_id, data={}) + area_id, TimeSeriesData(type=data_selection.type, name=area_id, data={}) ) - for var_index in area.variables: + for var_index, var in enumerate(area_df.variables): var = area_vars[var_index] - numerical_data = df.to_series(offset + var_index).cast(float).to_list() + numerical_data = df.to_series(var_index).cast(float).to_list() ts_data.data.setdefault(str(year), []).append( TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) ) diff --git a/antarest/output/storage/v2/iteration.py b/antarest/output/storage/v2/iteration.py new file mode 100644 index 0000000000..be8a58e327 --- /dev/null +++ b/antarest/output/storage/v2/iteration.py @@ -0,0 +1,98 @@ +# Copyright (c) 2026, 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 itertools +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Sequence + +from polars import DataFrame, col, scan_parquet +from polars.selectors import by_index + +from antarest.output.filestudy.model import VariableDescription +from antarest.output.storage.v2.dbmodel import ElementType +from antarest.output.storage.v2.metadata import IParquetOutputMetadata +from antarest.study.model import MatrixFrequency + + +def _parquet_file_name(element_type: ElementType, frequency: MatrixFrequency) -> str: + match element_type: + case "area": + return f"mc-ind_areas_{frequency.value}.parquet" + case _: + raise NotImplementedError("Not yet implemented") + + +def _parquet_file(parquet_dir: Path, element_type: ElementType, frequency: MatrixFrequency) -> Path: + return parquet_dir / _parquet_file_name(element_type, frequency) + + +@dataclass(frozen=True) +class AreaDataFrame: + year: int + area_id: str + variables: Sequence[VariableDescription] + data: DataFrame + + +# TODO: should go somewhere else +MC_IND_AREA_INDEX = ("mcYear", "area", "timeId") +MC_IND_AREA_COL_OFFSET = len(MC_IND_AREA_INDEX) + + +def iterate_areas_df( + output_metadata: IParquetOutputMetadata, + parquet_dir: Path, + frequency: MatrixFrequency, + years: Sequence[int], + areas: Sequence[str], + columns: Sequence[str], +) -> Iterable[AreaDataFrame]: + """ + Yields dataframes for the selected years and areas, in sorted order, years moving last. + + Implementation first scans the parquet file for the selected rows and columns, + then iterate on each couple year/area to yield the corresponding dataframe. + We take care of selecting, for each area, only the variables of that area. + """ + all_area_vars = output_metadata.get_variables("mc-ind", "area") + parquet_file = _parquet_file(parquet_dir, "area", frequency) + areas_df = scan_parquet(parquet_file) + if years: + areas_df = areas_df.filter(col("mcYear").is_in(years)) + if areas: + areas_df = areas_df.filter(col("area").is_in(areas)) + + selected_cols: set[int] = set() + if columns: + selected_cols = {c for c, v in enumerate(all_area_vars) if v.name in columns} + + actual_areas = output_metadata.mc_ind_areas + if areas: + actual_areas = [a for a in actual_areas if a.area_id in areas] + actual_areas = sorted(actual_areas, key=lambda a: a.area_id) + actual_years = years if years else output_metadata.mc_years + actual_years = sorted(actual_years) + + area_vars = {a.area_id: a for a in output_metadata.mc_ind_areas} + for year, area in itertools.product(actual_years, actual_areas): + vars_indices = area_vars[area.area_id].variables + if selected_cols: + vars_indices = [v for v in vars_indices if v in selected_cols] + vars = [all_area_vars[i] for i in vars_indices] + df = ( + areas_df.filter(col("area") == area.area_id) + .filter(col("mcYear") == year) + .select(by_index([MC_IND_AREA_COL_OFFSET + v for v in vars_indices])) + .collect() + ) + + yield AreaDataFrame(year=year, area_id=area.area_id, variables=vars, data=df) From 143ee0c63c70a7d7d192489f1b1b68f8677f301a Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 18:39:35 +0200 Subject: [PATCH 24/29] add some test, fix download impl Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/download.py | 7 +- antarest/output/storage/v2/iteration.py | 12 ++- tests/output/storage/v2/conftest.py | 57 ++++++++++ tests/output/storage/v2/test_download.py | 77 ++----------- tests/output/storage/v2/test_iteration.py | 126 ++++++++++++++++++++++ 5 files changed, 202 insertions(+), 77 deletions(-) create mode 100644 tests/output/storage/v2/conftest.py create mode 100644 tests/output/storage/v2/test_iteration.py diff --git a/antarest/output/storage/v2/download.py b/antarest/output/storage/v2/download.py index b8759e8708..99841408c3 100644 --- a/antarest/output/storage/v2/download.py +++ b/antarest/output/storage/v2/download.py @@ -32,8 +32,6 @@ def build_matrix_aggregation_result( output_metadata: IParquetOutputMetadata, parquet_dir: Path, data_selection: StudyDownloadDTO ) -> MatrixAggregationResultDTO: - area_vars = output_metadata.get_variables("mc-ind", "area") - element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system if data_selection.type in {StudyDownloadType.AREA, StudyDownloadType.DISTRICT}: @@ -46,12 +44,11 @@ def build_matrix_aggregation_result( data_selection.columns, ) for area_df in area_dfs: - year, area_id, df = area_df.year, area_df.area_id, area_df.data + year, area_id, df, vars = area_df.year, area_df.area_id, area_df.data, area_df.variables ts_data = element_results.setdefault( area_id, TimeSeriesData(type=data_selection.type, name=area_id, data={}) ) - for var_index, var in enumerate(area_df.variables): - var = area_vars[var_index] + for var_index, var in enumerate(vars): numerical_data = df.to_series(var_index).cast(float).to_list() ts_data.data.setdefault(str(year), []).append( TimeSerie(name=var.name, unit=var.unit_repr(), data=numerical_data) diff --git a/antarest/output/storage/v2/iteration.py b/antarest/output/storage/v2/iteration.py index be8a58e327..fd34f70270 100644 --- a/antarest/output/storage/v2/iteration.py +++ b/antarest/output/storage/v2/iteration.py @@ -59,12 +59,22 @@ def iterate_areas_df( """ Yields dataframes for the selected years and areas, in sorted order, years moving last. + Note that each area may have different variables and hence different dataframe shapes. + For example, if 2 areas have different thermal cluster groups, that will be the case. + Implementation first scans the parquet file for the selected rows and columns, then iterate on each couple year/area to yield the corresponding dataframe. We take care of selecting, for each area, only the variables of that area. + + Note on performance: + we scan from the parquet file for each area, which is likely sub-optimal. + Implementation may be tuned if considered useful later, for example by collecting a DataFrame + with all necessary data first (but that can cause out of memory errors). """ all_area_vars = output_metadata.get_variables("mc-ind", "area") parquet_file = _parquet_file(parquet_dir, "area", frequency) + + # Using polars lazy frame API to define the query into the underlying parquet file areas_df = scan_parquet(parquet_file) if years: areas_df = areas_df.filter(col("mcYear").is_in(years)) @@ -87,7 +97,6 @@ def iterate_areas_df( vars_indices = area_vars[area.area_id].variables if selected_cols: vars_indices = [v for v in vars_indices if v in selected_cols] - vars = [all_area_vars[i] for i in vars_indices] df = ( areas_df.filter(col("area") == area.area_id) .filter(col("mcYear") == year) @@ -95,4 +104,5 @@ def iterate_areas_df( .collect() ) + vars = [all_area_vars[i] for i in vars_indices] yield AreaDataFrame(year=year, area_id=area.area_id, variables=vars, data=df) diff --git a/tests/output/storage/v2/conftest.py b/tests/output/storage/v2/conftest.py new file mode 100644 index 0000000000..951d6a7d74 --- /dev/null +++ b/tests/output/storage/v2/conftest.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026, 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 pathlib import Path + +import pytest +from sqlalchemy.orm import Session + +from antarest.output.filestudy.model import FileOutput +from antarest.output.storage.v2.dbmodel import DbParquetOutput +from antarest.output.storage.v2.metadata import IParquetOutputMetadata, ParquetOuputMetadataImpl +from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database +from antarest.output.storage.v2.variables_storage import create_parquet_files + + +@pytest.fixture +def output_dir(data_dir: Path) -> Path: + return data_dir / "20260810-1420eco-thermal_groups" + + +@pytest.fixture +def parquet_dir(tmp_path: Path) -> Path: + dir = tmp_path / "output" + dir.mkdir() + return dir + + +@pytest.fixture +def parquet_metadata(output_dir: Path, parquet_dir: Path, db_session: Session) -> IParquetOutputMetadata: + """ + Imports 20260810-1420eco-thermal_groups to parquet and return the associated metadata + """ + + db_output = DbParquetOutput(id=0, mc_years=[1, 2]) + db_session.add(db_output) + db_session.flush() + + file_output = FileOutput(output_dir) + extract_output_variables_to_database(db_session, db_output.id, file_output) + db_session.flush() + + output_metadata = ParquetOuputMetadataImpl(db_session, db_output.id) + create_parquet_files(output_metadata, file_output, parquet_dir) + + assert len(list(parquet_dir.iterdir())) == 1 + monthly_file = parquet_dir / "mc-ind_areas_monthly.parquet" + assert monthly_file.is_file() + + return output_metadata diff --git a/tests/output/storage/v2/test_download.py b/tests/output/storage/v2/test_download.py index e2de59c113..5a55d53ce8 100644 --- a/tests/output/storage/v2/test_download.py +++ b/tests/output/storage/v2/test_download.py @@ -11,57 +11,18 @@ # This file is part of the Antares project. from pathlib import Path -import pytest -from sqlalchemy.orm import Session - -from antarest.output.filestudy.model import FileOutput from antarest.output.model import StudyDownloadDTO, StudyDownloadType -from antarest.output.storage.v2.dbmodel import DbParquetOutput from antarest.output.storage.v2.download import ( build_matrix_aggregation_result, ) -from antarest.output.storage.v2.metadata import ParquetOuputMetadataImpl -from antarest.output.storage.v2.variables_fetching import get_variables_index -from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database -from antarest.output.storage.v2.variables_storage import extract_areas_refacto +from antarest.output.storage.v2.metadata import IParquetOutputMetadata from antarest.study.model import MatrixFrequency -@pytest.fixture -def output_dir(data_dir: Path) -> Path: - return data_dir / "20260810-1420eco-thermal_groups" - - -def test_download_areas(output_dir: Path, db_session: Session, tmp_path: Path) -> None: - # TODO: simplify setup - - # Setup - - db_output = DbParquetOutput(id=0, playlist=[1, 2]) - db_session.add(db_output) - db_session.flush() - - file_output = FileOutput(output_dir) - extract_output_variables_to_database(db_session, db_output.id, file_output) - db_session.flush() - - extract_output_variables_to_database(db_session, 0, file_output) - index = get_variables_index(db_session, 0) - - target_dir = tmp_path / "output" - target_dir.mkdir() - - extract_areas_refacto(index, file_output, target_dir) - - assert len(list(target_dir.iterdir())) == 1 - monthly_file = target_dir / "mc-ind_areas_monthly.parquet" - assert monthly_file.is_file() +def test_download_areas(parquet_dir: Path, parquet_metadata: IParquetOutputMetadata) -> None: - # Actual test - - output_metadata = ParquetOuputMetadataImpl(db_session, db_output.id) data_selection = StudyDownloadDTO(type=StudyDownloadType.AREA, years=[], level=MatrixFrequency.MONTHLY, filter=[]) - aggregate = build_matrix_aggregation_result(output_metadata, target_dir, data_selection) + aggregate = build_matrix_aggregation_result(parquet_metadata, parquet_dir, data_selection) year1_st_by_area = {data.name: data.data["1"] for data in aggregate.data} es_variables = [ts.name for ts in year1_st_by_area["es"]] @@ -90,38 +51,12 @@ def test_download_areas(output_dir: Path, db_session: Session, tmp_path: Path) - ] -def test_download_district(output_dir: Path, db_session: Session, tmp_path: Path) -> None: - # TODO: simplify setup - # TODO: make it pass - - # Setup - - db_output = DbParquetOutput(id=0, playlist=[1, 2]) - db_session.add(db_output) - db_session.flush() - - file_output = FileOutput(output_dir) - extract_output_variables_to_database(db_session, db_output.id, file_output) - db_session.flush() - - extract_output_variables_to_database(db_session, 0, file_output) - index = get_variables_index(db_session, 0) - - target_dir = tmp_path / "output" - target_dir.mkdir() - - extract_areas_refacto(index, file_output, target_dir) - - assert len(list(target_dir.iterdir())) == 1 - monthly_file = target_dir / "mc-ind_areas_monthly.parquet" - assert monthly_file.is_file() - - # Actual test +def test_download_district(parquet_dir: Path, parquet_metadata: IParquetOutputMetadata) -> None: - output_metadata = ParquetOuputMetadataImpl(db_session, db_output.id) + # Like in filesystem implementation, DISTRICT is handled identically to AREA ... data_selection = StudyDownloadDTO(type=StudyDownloadType.DISTRICT, years=[1], level=MatrixFrequency.MONTHLY) - aggregate = build_matrix_aggregation_result(output_metadata, target_dir, data_selection) + aggregate = build_matrix_aggregation_result(parquet_metadata, parquet_dir, data_selection) year1_st_by_area = {data.name: data.data["1"] for data in aggregate.data} all_areas_variables = [ts.name for ts in year1_st_by_area["@ all areas"]] diff --git a/tests/output/storage/v2/test_iteration.py b/tests/output/storage/v2/test_iteration.py new file mode 100644 index 0000000000..fba2bc6f0c --- /dev/null +++ b/tests/output/storage/v2/test_iteration.py @@ -0,0 +1,126 @@ +# Copyright (c) 2026, 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 pathlib import Path + +from antarest.output.filestudy.model import VariableDescription +from antarest.output.storage.v2.iteration import iterate_areas_df +from antarest.output.storage.v2.metadata import IParquetOutputMetadata +from antarest.study.model import MatrixFrequency + + +def test_iterate_areas(parquet_dir: Path, parquet_metadata: IParquetOutputMetadata) -> None: + + area_dfs = list(iterate_areas_df(parquet_metadata, parquet_dir, MatrixFrequency.MONTHLY, [], [], [])) + + assert len(area_dfs) == 6 # 2 years, 3 "areas" (1 district ...) + + assert [(a.year, a.area_id) for a in area_dfs] == [ + (1, "@ all areas"), + (1, "es"), + (1, "fr"), + (2, "@ all areas"), + (2, "es"), + (2, "fr"), + ] + + # Check district variables + assert area_dfs[0].variables == [ + VariableDescription(name="CO2 EMIS.", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="ES_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), + VariableDescription(name="FR_NUCLEAR_TH_PROD", unit="MWh", statistic_type=None), + ] + + # check we get ES variables for ES, FR variables for FR + assert area_dfs[1].variables == [ + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), + VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + ] + + assert area_dfs[2].variables == [ + VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), + VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None), + VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), + VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), + VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), + VariableDescription(name="NP COST", unit="Euro", statistic_type=None), + VariableDescription(name="NODU", unit=None, statistic_type=None), + VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), + ] + + # Check content for one column + assert area_dfs[2].data.to_series(1).to_list() == [ + 595200.0, + 537600.0, + 595200.0, + 576000.0, + 595200.0, + 576000.0, + 595200.0, + 595200.0, + 576000.0, + 595200.0, + 576000.0, + 576000.0, + ] + + +def test_iterate_areas_filters(parquet_dir: Path, parquet_metadata: IParquetOutputMetadata) -> None: + + area_dfs = list(iterate_areas_df(parquet_metadata, parquet_dir, MatrixFrequency.MONTHLY, [1], [], [])) + + assert [(a.year, a.area_id) for a in area_dfs] == [ + (1, "@ all areas"), + (1, "es"), + (1, "fr"), + ] + + area_dfs = list(iterate_areas_df(parquet_metadata, parquet_dir, MatrixFrequency.MONTHLY, [1], ["fr"], [])) + + assert [(a.year, a.area_id) for a in area_dfs] == [ + (1, "fr"), + ] + + area_dfs = list( + iterate_areas_df(parquet_metadata, parquet_dir, MatrixFrequency.MONTHLY, [1], ["fr"], ["FR_NUCLEAR"]) + ) + + assert [(a.year, a.area_id) for a in area_dfs] == [ + (1, "fr"), + ] + assert area_dfs[0].variables == [VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None)] + assert area_dfs[0].data.columns == ["FR_NUCLEAR__MWh"] + assert area_dfs[0].data.to_series(0).to_list() == [ + 595200.0, + 537600.0, + 595200.0, + 576000.0, + 595200.0, + 576000.0, + 595200.0, + 595200.0, + 576000.0, + 595200.0, + 576000.0, + 576000.0, + ] From 3bb9d3a70ccbd87093577976ffc72255ea4aaea9 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 19:06:19 +0200 Subject: [PATCH 25/29] remove unused classes Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/repository.py | 8 ----- .../output/storage/v2/variables_fetching.py | 30 ------------------- 2 files changed, 38 deletions(-) diff --git a/antarest/output/storage/v2/repository.py b/antarest/output/storage/v2/repository.py index cd79fd4d56..b0dcccf4d7 100644 --- a/antarest/output/storage/v2/repository.py +++ b/antarest/output/storage/v2/repository.py @@ -10,7 +10,6 @@ # # This file is part of the Antares project. from collections.abc import Iterator -from typing import Iterable from sqlalchemy import Boolean, Column, ForeignKeyConstraint, Integer, String, Table, delete, select from sqlalchemy.orm import Mapped, Session, mapped_column @@ -20,7 +19,6 @@ from antarest.core.utils.sql_utils import upsert_one from antarest.launcher.model import LogType from antarest.output.model import OutputVariablesList -from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable class DbOutputMetadataV2(Base): @@ -176,9 +174,3 @@ def save_output_variables_list(self, study_id: str, output_id: str, variables_li db_model = DbOutputVariablesV2.from_model(study_id, output_id, variables_list) self.session.add(db_model) self.session.commit() - - def save_variables(self, variables: Iterable[DbParquetVariable]) -> None: - self.session.add_all(variables) - - def save_areas(self, areas: Iterable[DbParquetArea]) -> None: - self.session.add_all(areas) diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index 91dfc33eef..928ed29224 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -14,7 +14,6 @@ Fetching variables metadata from the database """ -from dataclasses import dataclass from typing import Iterable, Sequence from sqlalchemy import select @@ -24,25 +23,6 @@ from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ElementType, ScenarioAggregation -@dataclass(frozen=True) -class VariableColumn: - """ - Attributes: - column: offset of the column in the parquet file (actual col will be number of index cols + this offset) - """ - - column: int - name: str - unit: str | None - statistic_type: str | None - - -@dataclass(frozen=True) -class AreaVariables: - area_id: str - variables: Sequence[VariableColumn] - - class VariablesIndex: """ Helper class to retrieve variable info from DB models. @@ -67,16 +47,6 @@ def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementT """ return [_to_var_desc(v) for v in self._get_db_vars(aggregation, element_type)] - def get_variable_columns(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableColumn]: - """ - Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) - """ - return [_to_var_col(v) for v in self._get_db_vars(aggregation, element_type)] - - -def _to_var_col(db_var: DbParquetVariable) -> VariableColumn: - return VariableColumn(db_var.column, db_var.name, db_var.unit, db_var.statistic_type) - def _to_var_desc(db_var: DbParquetVariable) -> VariableDescription: return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) From a0b51194abcf07855e7af83ded9d88bfd2dfe4cb Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 19:08:26 +0200 Subject: [PATCH 26/29] remove unused methods Signed-off-by: Sylvain Leclerc --- .../output/storage/v2/variables_fetching.py | 22 ------------------- 1 file changed, 22 deletions(-) diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py index 928ed29224..720c4d4567 100644 --- a/antarest/output/storage/v2/variables_fetching.py +++ b/antarest/output/storage/v2/variables_fetching.py @@ -50,25 +50,3 @@ def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementT def _to_var_desc(db_var: DbParquetVariable) -> VariableDescription: return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) - - -def get_variables_index(session: Session, output_id: int) -> VariablesIndex: - output_variables = session.execute( - select(DbParquetVariable).where(DbParquetVariable.output_id == output_id) - ).scalars() - return VariablesIndex(output_variables) - - -def get_area_variables( - session: Session, output_id: int, aggregation: ScenarioAggregation, area_id: str -) -> list[VariableDescription]: - - # All variables, should load fast ? - variables_index = get_variables_index(session, output_id) - - # Get area information - area = session.execute(select(DbParquetArea).where(DbParquetArea.area_id == area_id)).scalar_one() - - all_areas_vars = variables_index.get_variables(aggregation, "area") - cols = area.mc_all_vars if aggregation == "mc-all" else area.mc_ind_vars - return [all_areas_vars[c] for c in cols] From 87ce1b8ce298f5f5933d9313a6aa411d58dbeffd Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 19:10:18 +0200 Subject: [PATCH 27/29] move helper class Signed-off-by: Sylvain Leclerc --- antarest/output/storage/v2/metadata.py | 32 +++++++++++- .../output/storage/v2/variables_fetching.py | 52 ------------------- 2 files changed, 30 insertions(+), 54 deletions(-) delete mode 100644 antarest/output/storage/v2/variables_fetching.py diff --git a/antarest/output/storage/v2/metadata.py b/antarest/output/storage/v2/metadata.py index 4e7cb7602b..eb2040c02b 100644 --- a/antarest/output/storage/v2/metadata.py +++ b/antarest/output/storage/v2/metadata.py @@ -12,7 +12,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from functools import cached_property -from typing import Sequence +from typing import Iterable, Sequence from sqlalchemy import select from sqlalchemy.orm import Session @@ -27,7 +27,6 @@ ElementType, ScenarioAggregation, ) -from antarest.output.storage.v2.variables_fetching import VariablesIndex from antarest.study.model import MatrixFrequency @@ -82,6 +81,35 @@ def mc_ind_areas(self) -> Sequence[AreaVariables]: """ +class VariablesIndex: + """ + Helper class to retrieve variable info from DB models. + """ + + def __init__(self, variables: Iterable[DbParquetVariable]) -> None: + vars: dict[tuple[ScenarioAggregation, ElementType], list[DbParquetVariable]] = {} + for v in variables: + vars.setdefault((v.scenario_aggregation, v.element_type), []).append(v) + + self._variables = {k: sorted(v, key=lambda v: v.column) for k, v in vars.items()} # sort by columns + + def _get_db_vars(self, aggregation: ScenarioAggregation, element_type: ElementType) -> Sequence[DbParquetVariable]: + """ + Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) + """ + return self._variables.get((aggregation, element_type), []) + + def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: + """ + Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) + """ + return [_to_var_desc(v) for v in self._get_db_vars(aggregation, element_type)] + + +def _to_var_desc(db_var: DbParquetVariable) -> VariableDescription: + return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) + + class ParquetOuputMetadataImpl(IParquetOutputMetadata): """ Implementation which gets metadata from the DB, as lazily as possible. diff --git a/antarest/output/storage/v2/variables_fetching.py b/antarest/output/storage/v2/variables_fetching.py deleted file mode 100644 index 720c4d4567..0000000000 --- a/antarest/output/storage/v2/variables_fetching.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) 2026, 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. - -""" -Fetching variables metadata from the database -""" - -from typing import Iterable, Sequence - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from antarest.output.filestudy.model import VariableDescription -from antarest.output.storage.v2.dbmodel import DbParquetArea, DbParquetVariable, ElementType, ScenarioAggregation - - -class VariablesIndex: - """ - Helper class to retrieve variable info from DB models. - """ - - def __init__(self, variables: Iterable[DbParquetVariable]) -> None: - vars: dict[tuple[ScenarioAggregation, ElementType], list[DbParquetVariable]] = {} - for v in variables: - vars.setdefault((v.scenario_aggregation, v.element_type), []).append(v) - - self._variables = {k: sorted(v, key=lambda v: v.column) for k, v in vars.items()} # sort by columns - - def _get_db_vars(self, aggregation: ScenarioAggregation, element_type: ElementType) -> Sequence[DbParquetVariable]: - """ - Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) - """ - return self._variables.get((aggregation, element_type), []) - - def get_variables(self, aggregation: ScenarioAggregation, element_type: ElementType) -> list[VariableDescription]: - """ - Get all variables for the specified "mc-ind/mc-all" and element type (areas, links, ...) - """ - return [_to_var_desc(v) for v in self._get_db_vars(aggregation, element_type)] - - -def _to_var_desc(db_var: DbParquetVariable) -> VariableDescription: - return VariableDescription(db_var.name, db_var.unit, db_var.statistic_type) From abd9ae30b105ce0fe69206f1a08c627677f1b481 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 19:16:12 +0200 Subject: [PATCH 28/29] remove obsolete unit test Signed-off-by: Sylvain Leclerc --- .../storage/v2/test_variables_fetching.py | 93 ------------------- 1 file changed, 93 deletions(-) delete mode 100644 tests/output/storage/v2/test_variables_fetching.py diff --git a/tests/output/storage/v2/test_variables_fetching.py b/tests/output/storage/v2/test_variables_fetching.py deleted file mode 100644 index 2fbe8be389..0000000000 --- a/tests/output/storage/v2/test_variables_fetching.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) 2026, 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 pathlib import Path - -import pytest -from sqlalchemy.orm import Session - -from antarest.output.filestudy.model import FileOutput, VariableDescription -from antarest.output.storage.v2.dbmodel import DbParquetOutput -from antarest.output.storage.v2.variables_fetching import get_area_variables -from antarest.output.storage.v2.variables_parsing import extract_output_variables_to_database - - -@pytest.fixture -def output_dir(data_dir: Path) -> Path: - return data_dir / "20260810-1420eco-thermal_groups" - - -def test_get_area_variables(db_session: Session, output_dir: Path) -> None: - - db_output = DbParquetOutput(id=0) - db_session.add(db_output) - db_session.flush() - - file_output = FileOutput(output_dir) - extract_output_variables_to_database(db_session, db_output.id, file_output) - db_session.flush() - - assert get_area_variables(db_session, db_output.id, "mc-ind", "fr") == [ - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), - VariableDescription(name="FR_NUCLEAR", unit="MWh", statistic_type=None), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), - VariableDescription(name="NP COST", unit="Euro", statistic_type=None), - VariableDescription(name="NODU", unit=None, statistic_type=None), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), - ] - - assert get_area_variables(db_session, db_output.id, "mc-ind", "es") == [ - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type=None), - VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type=None), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type=None), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type=None), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type=None), - VariableDescription(name="NP COST", unit="Euro", statistic_type=None), - VariableDescription(name="NODU", unit=None, statistic_type=None), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type=None), - ] - - assert get_area_variables(db_session, db_output.id, "mc-all", "es") == [ - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="EXP"), - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="std"), - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="min"), - VariableDescription(name="CO2 EMIS.", unit="Tons", statistic_type="max"), - VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="EXP"), - VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="std"), - VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="min"), - VariableDescription(name="ES_NUCLEAR", unit="MWh", statistic_type="max"), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type="EXP"), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type="std"), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type="min"), - VariableDescription(name="AVL DTG", unit="MWh", statistic_type="max"), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type="EXP"), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type="std"), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type="min"), - VariableDescription(name="DTG MRG", unit="MWh", statistic_type="max"), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type="EXP"), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type="std"), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type="min"), - VariableDescription(name="MAX MRG", unit="MWh", statistic_type="max"), - VariableDescription(name="NP COST", unit="Euro", statistic_type="EXP"), - VariableDescription(name="NP COST", unit="Euro", statistic_type="std"), - VariableDescription(name="NP COST", unit="Euro", statistic_type="min"), - VariableDescription(name="NP COST", unit="Euro", statistic_type="max"), - VariableDescription(name="NODU", unit=None, statistic_type="EXP"), - VariableDescription(name="NODU", unit=None, statistic_type="std"), - VariableDescription(name="NODU", unit=None, statistic_type="min"), - VariableDescription(name="NODU", unit=None, statistic_type="max"), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type="EXP"), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type="std"), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type="min"), - VariableDescription(name="RES LOAD", unit="MWh", statistic_type="max"), - ] From b9eed3dd6cb37034201a441f5446660e503bb933 Mon Sep 17 00:00:00 2001 From: Sylvain Leclerc Date: Tue, 25 Aug 2026 19:36:17 +0200 Subject: [PATCH 29/29] remove unused fixture Signed-off-by: Sylvain Leclerc --- tests/conftest.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 2d00c9b426..5e5d0e1a3e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -68,11 +68,6 @@ def project_path() -> Path: return PROJECT_DIR -@pytest.fixture(scope="session") -def test_root() -> Path: - return HERE - - @pytest.fixture def ini_cleaner() -> Callable[[str], str]: def cleaner(txt: str) -> str: