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/filestudy/model.py b/antarest/output/filestudy/model.py index d30f64d1a6..4322f3b9bd 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,117 @@ 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 mc_ind_link_ids(self) -> tuple[str, ...]: + """ + 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 mc_ind_area_ids(self) -> tuple[str, ...]: + """ + 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. + """ + 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, 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 + + 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..a93680daa6 --- /dev/null +++ b/antarest/output/storage/v2/dbmodel.py @@ -0,0 +1,108 @@ +# 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, Literal, TypeAlias + +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 + +ElementType: TypeAlias = Literal[ + "area", + "link", + "binding_constraint", + "thermal_cluster", + "renewable_cluster", + "short_term_storage", +] + +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 + # study_id / output_id + + __tablename__ = "parquet_output" + + id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True) + mc_years: Mapped[list[int]] = mapped_column(IntList) + + +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 + + 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 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 + + 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(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..99841408c3 --- /dev/null +++ b/antarest/output/storage/v2/download.py @@ -0,0 +1,60 @@ +# 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 +""" + +from pathlib import Path + +from antarest.output.model import ( + MatrixAggregationResultDTO, + StudyDownloadDTO, + StudyDownloadType, + TimeSerie, + TimeSeriesData, +) +from antarest.output.storage.v2.iteration import iterate_areas_df +from antarest.output.storage.v2.metadata import IParquetOutputMetadata + + +def build_matrix_aggregation_result( + output_metadata: IParquetOutputMetadata, parquet_dir: Path, data_selection: StudyDownloadDTO +) -> MatrixAggregationResultDTO: + + element_results: dict[str, TimeSeriesData] = {} # one TimeSeriesData for each element of the system + + 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, 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(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) + ) + + return MatrixAggregationResultDTO( + index=output_metadata.get_time_index(data_selection.level), + data=list(element_results.values()), + ) diff --git a/antarest/output/storage/v2/iteration.py b/antarest/output/storage/v2/iteration.py new file mode 100644 index 0000000000..fd34f70270 --- /dev/null +++ b/antarest/output/storage/v2/iteration.py @@ -0,0 +1,108 @@ +# 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. + + 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)) + 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] + 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() + ) + + 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/antarest/output/storage/v2/metadata.py b/antarest/output/storage/v2/metadata.py new file mode 100644 index 0000000000..eb2040c02b --- /dev/null +++ b/antarest/output/storage/v2/metadata.py @@ -0,0 +1,167 @@ +# 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 Iterable, 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.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 ... + + 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 + @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 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. + + 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.mc_years + + @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/storage.py b/antarest/output/storage/v2/storage.py index c2364db47a..e331de3ae8 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,15 @@ 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, 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,9 +227,6 @@ def import_output( simulation_range = _extract_simulation_range(dir_path) - 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, @@ -243,6 +245,14 @@ def import_output( ) ) + 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) @@ -449,4 +459,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) diff --git a/antarest/output/storage/v2/variables_parsing.py b/antarest/output/storage/v2/variables_parsing.py new file mode 100644 index 0000000000..e6ce0472bd --- /dev/null +++ b/antarest/output/storage/v2/variables_parsing.py @@ -0,0 +1,151 @@ +# 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. + +""" +Parsing of variables metadata from file studies, in order to populate the database +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +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, ElementType, 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]] = {} + + # data source depends on aggregation type + get_file: Callable[[str, MatrixFrequency], Path | None] + area_ids: Iterable[str] + 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 + ) + + 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) + + 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): + 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 _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. + """ + 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 + 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 + 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) diff --git a/antarest/output/storage/v2/variables_storage.py b/antarest/output/storage/v2/variables_storage.py index 4b8e1dcb3a..733c273e8b 100644 --- a/antarest/output/storage/v2/variables_storage.py +++ b/antarest/output/storage/v2/variables_storage.py @@ -22,30 +22,40 @@ import shutil import tempfile from collections.abc import Iterator, Sequence +from dataclasses import dataclass from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias import polars as pl +import polars.selectors as pls +import pyarrow as pa +from polars import Float64 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, ) 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, TIME_ID_COL, + FileOutput, MCAllAreasQueryFile, MCAllLinksQueryFile, MCIndAreasQueryFile, MCIndLinksQueryFile, MCRoot, QueryFileType, + VariableDescription, find_mode_dir, get_output_object_type, ) +from antarest.output.storage.v2.metadata import IParquetOutputMetadata from antarest.study.model import MatrixFrequency logger = logging.getLogger(__name__) @@ -392,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) 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/conftest.py b/tests/output/conftest.py new file mode 100644 index 0000000000..c589615d0d --- /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() -> 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/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/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/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 new file mode 100644 index 0000000000..5a55d53ce8 --- /dev/null +++ b/tests/output/storage/v2/test_download.py @@ -0,0 +1,74 @@ +# 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.model import StudyDownloadDTO, StudyDownloadType +from antarest.output.storage.v2.download import ( + build_matrix_aggregation_result, +) +from antarest.output.storage.v2.metadata import IParquetOutputMetadata +from antarest.study.model import MatrixFrequency + + +def test_download_areas(parquet_dir: Path, parquet_metadata: IParquetOutputMetadata) -> None: + + data_selection = StudyDownloadDTO(type=StudyDownloadType.AREA, years=[], level=MatrixFrequency.MONTHLY, filter=[]) + 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"]] + 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(parquet_dir: Path, parquet_metadata: IParquetOutputMetadata) -> None: + + # 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(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"]] + + 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", + ] 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, + ] 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..b7e0d1d216 --- /dev/null +++ b/tests/output/storage/v2/test_variable_storage.py @@ -0,0 +1,114 @@ +# 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 polars as pl +import pytest +from polars.testing import assert_frame_equal +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_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: + var_cols = [ + VariableDescription("var1", None, "exp"), + VariableDescription("var2", None, None), + VariableDescription("var3", None, None), + ] + 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()), + ] + ), + index_cols=["area", "timeId"], + var_cols=[VariableDescription("var3", None, None), VariableDescription("var1", None, "exp")], + ) + 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="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()), + ] + ) + + 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 ... + + 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) + + assert len(list(target_dir.iterdir())) == 1 + 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", + ] 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..16ba8d1cdc --- /dev/null +++ b/tests/output/storage/v2/test_variables_parsing.py @@ -0,0 +1,203 @@ +# 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_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, + ], + }