utils/rfsd_loader.py - master - Zerg/Intermediate_certification2

utils/rfsd_loader.py
370 строк10 KB
zerg
Пофиксил файлы
11 авг 2026, 00:21
11 авг 2026, 00:217c22706
from __future__ import annotations

import logging
from collections import defaultdict
from pathlib import Path

import polars as pl
from huggingface_hub import snapshot_download

logger = logging.getLogger(__name__)

# ссылка на файл с человеческими именами столбцов для переименования
DEFAULT_RENAME_URL = (
    "https://raw.githubusercontent.com/irlcode/RFSD/main/"
    "aux_data/descriptive_names_dict.csv"
)


class RFSDLoader:
    """
    Загрузка и предварительная обработка датасета RFSD.
    Все операции построены на Lazy API Polars.
    """

    def __init__(
        self,
        dataset_dir: str | Path,
        output_dir: str | Path,
        rename_source: str | Path = DEFAULT_RENAME_URL,
    ) -> None:

        self.dataset_dir = Path(dataset_dir)
        self.output_dir = Path(output_dir)
        self.rename_source = rename_source

        self.output_dir.mkdir(
            parents=True,
            exist_ok=True,
        )

        self._rename_map: dict[str, str] | None = None

    # ------------------------------------------------------------------
    # Download
    # ------------------------------------------------------------------

    def download(
        self,
        repo_id: str = "irlspbru/RFSD",
        workers: int = 8,
    ) -> str:
        """
        Скачать датасет с HuggingFace.
        """

        logger.info("Downloading RFSD...")

        path = snapshot_download(
            repo_id=repo_id,
            repo_type="dataset",
            local_dir=str(self.dataset_dir),
            max_workers=workers,
        )

        logger.info("Download completed.")

        return path

    # ------------------------------------------------------------------
    # Rename map
    # ------------------------------------------------------------------

    @property
    def rename_map(self) -> dict[str, str]:
        """
        Ленивая загрузка карты переименования.
        """

        if self._rename_map is None:
            self._rename_map = self._load_rename_map()

        return self._rename_map

    def _load_rename_map(self) -> dict[str, str]:
        """
        Загрузить карту переименования колонок.

        Если descriptive-названия повторяются,
        автоматически добавить суффиксы:

            B_fin_invest
            B_fin_invest_2
            B_fin_invest_3
        """

        source = Path(self.rename_source)

        if source.exists():
            df = pl.read_csv(source)
        else:
            df = pl.read_csv(self.rename_source)

        return self._make_unique_names(df)

    @staticmethod
    def _make_unique_names(
        rename_df: pl.DataFrame,
    ) -> dict[str, str]:
        """
        Сделать descriptive-названия уникальными.
        """

        used = defaultdict(int)
        mapping: dict[str, str] = {}

        originals = rename_df["original"].to_list()
        descriptives = rename_df["descriptive"].to_list()

        for original, descriptive in zip(originals, descriptives):
            if original is None or descriptive is None:
                continue

            used[descriptive] += 1

            if used[descriptive] == 1:
                mapping[original] = descriptive
            else:
                mapping[original] = f"{descriptive}_{used[descriptive]}"

        logger.info(
            "Loaded %d column mappings.",
            len(mapping),
        )

        return mapping

    # ------------------------------------------------------------------
    # Dataset discovery
    # ------------------------------------------------------------------

    def year_directories(self) -> list[Path]:
        """
        Найти только реальные каталоги year=XXXX.
        """

        root = self.dataset_dir / "RFSD"

        if not root.exists():
            root = self.dataset_dir

        return sorted(p for p in root.glob("year=*") if p.is_dir())

    # ------------------------------------------------------------------
    # Scan
    # ------------------------------------------------------------------

    @staticmethod
    def scan_year(
        year_dir: Path,
    ) -> pl.LazyFrame:
        """
        Лениво открыть датасет одного года.
        """

        parquet_files = list(year_dir.glob("*.parquet"))

        if not parquet_files:
            raise FileNotFoundError(f"Parquet files not found: {year_dir}")

        return pl.scan_parquet(
            str(year_dir / "*.parquet"),
            missing_columns="insert",
        )

        # ------------------------------------------------------------------

    # Processing
    # ------------------------------------------------------------------

    def process_year(
        self,
        year_dir: Path,
        rename_columns: bool = True,
    ) -> Path:
        """
        Обработать датасет за один год.

        Parameters
        ----------
        year_dir
            Каталог year=XXXX.
        rename_columns
            Переименовывать ли колонки.

        Returns
        -------
        Path
            Путь к сохраненному parquet.
        """

        year = year_dir.name.split("=")[1]

        logger.info("Processing %s...", year)

        lf = self.scan_year(year_dir)

        if rename_columns:
            schema = lf.collect_schema()

            rename_dict = {
                old: new
                for old, new in self.rename_map.items()
                if old in schema.names()
            }

            if rename_dict:
                lf = lf.rename(rename_dict)

                logger.info(
                    "%s: renamed %d columns",
                    year,
                    len(rename_dict),
                )

        output = self.output_dir / f"{year}.parquet"

        lf.sink_parquet(output)

        logger.info("%s saved.", output.name)

        return output

    # ------------------------------------------------------------------

    def process_all_years(
        self,
        start_year: int | None = None,
        end_year: int | None = None,
        rename_columns: bool = True,
    ) -> list[Path]:
        """
        Обработать все найденные годы.
        """

        outputs: list[Path] = []

        for year_dir in self.year_directories():
            year = int(year_dir.name.split("=")[1])

            if start_year is not None and year < start_year:
                continue

            if end_year is not None and year > end_year:
                continue

            try:
                output = self.process_year(
                    year_dir,
                    rename_columns=rename_columns,
                )

                outputs.append(output)

            except Exception:
                logger.exception(
                    "Failed processing year %s",
                    year,
                )

        return outputs

    # ------------------------------------------------------------------
    # Schema
    # ------------------------------------------------------------------

    def compare_schemas(self) -> dict[str, list[str]]:
        """
        Сравнить схемы обработанных parquet-файлов.

        Returns
        -------
        dict
            Год -> отсутствующие столбцы.
        """

        schemas: dict[str, set[str]] = {}

        for file in sorted(self.output_dir.glob("*.parquet")):
            schema = pl.scan_parquet(file).collect_schema().names()

            schemas[file.stem] = set(schema)

        if not schemas:
            logger.warning("No parquet files found.")

            return {}

        all_columns = set.union(*schemas.values())

        differences: dict[str, list[str]] = {}

        for year, columns in schemas.items():
            missing = sorted(all_columns - columns)

            if missing:
                differences[year] = missing

                logger.info(
                    "%s missing %d columns",
                    year,
                    len(missing),
                )

        return differences

    # ------------------------------------------------------------------
    # Merge
    # ------------------------------------------------------------------

    def merge_years(
        self,
        start_year: int,
        end_year: int,
        output_name: str | None = None,
    ) -> Path:
        """
        Объединить несколько parquet-файлов.

        Parameters
        ----------
        start_year
            Первый год.

        end_year
            Последний год.

        output_name
            Имя результирующего файла.
        """

        files = []

        for year in range(start_year, end_year + 1):
            file = self.output_dir / f"{year}.parquet"

            if not file.exists():
                raise FileNotFoundError(file)

            files.append(file)

        logger.info(
            "Merging %d datasets...",
            len(files),
        )

        lf = pl.concat(
            [pl.scan_parquet(file) for file in files],
            how="diagonal_relaxed",
        )

        if output_name is None:
            output_name = f"rfsd_{start_year}_{end_year}.parquet"

        output = self.output_dir / output_name

        lf.sink_parquet(output)

        logger.info("Merged dataset saved.")

        return output