utils/preprocessing.py - master - Zerg/Intermediate_certification2

utils/preprocessing.py
138 строк4 KB
zerg
Создан preprocessing.py
8 часов назад
8 часов назадd80b875
from pathlib import Path

import polars as pl

from utils.rfsd_utils import PROCESSED_DIR


def drop_columns(
    lazy_frame: pl.LazyFrame,
    missing_threshold_percent: float = 5.0,
) -> pl.LazyFrame:
    """Удалztn столбцы, в которых доля пропусков выше заданного порога."""

    total_rows = lazy_frame.select(pl.len()).collect(engine="streaming").item()

    missing_stats = (
        lazy_frame.select(pl.all().null_count())
        .collect(engine="streaming")
        .unpivot(
            variable_name="column",
            value_name="count_missing",
        )
        .with_columns(
            (pl.col("count_missing") / total_rows * 100).alias("count_missing_percent")
        )
    )

    cols_to_drop = (
        missing_stats.filter(
            pl.col("count_missing_percent") > missing_threshold_percent
        )
        .get_column("column")
        .to_list()
    )

    return lazy_frame.drop(cols_to_drop)


def clean_dataset(
    lazy_frame: pl.LazyFrame,
    output_path: str | None = None,
    missing_threshold_percent: float = 5.0,
) -> pl.LazyFrame:
    """
    Очищает датасет.

    Этапы:
        1. Удалить столбцы, у которых доля пропусков превышает порог.
        2. Удалить оставшиеся строки с пропусками.
        3. При необходимости сохранить результат в Parquet.
    """

    cleaned_columns = drop_columns(
        lazy_frame,
        missing_threshold_percent=missing_threshold_percent,
    )
    cleaned = cleaned_columns.drop_nulls()

    if output_path is not None:
        cleaned.sink_parquet(output_path)
        print(f"Сохранено: {output_path}")

    result_rows = cleaned.select(pl.len()).collect(engine="streaming").item()
    result_columns = len(cleaned.collect_schema())

    print(
        f"Форма итогового датасета: {result_rows:,} строк × {result_columns} столбцов"
    )

    return cleaned


def get_missing_info(
    df: pl.LazyFrame,
    metric: str = "count",  # "count" or "percent"
) -> pl.DataFrame:
    """
    Расчет информации по пропускам в данных
    """
    missing_stats = (
        df.select(pl.all().null_count())
        .collect(engine="streaming")
        .unpivot(
            variable_name="column",
            value_name="count_missing",
        )
    )

    if metric == "percent":
        total_rows = df.select(pl.len()).collect(engine="streaming").item()
        return missing_stats.with_columns(
            (pl.col("count_missing") / total_rows * 100).alias("count_missing_percent")
        ).sort("count_missing_percent", descending=True)

    return missing_stats.sort("count_missing", descending=True)


def summarize_missing_parquet_datasets(
    parquet_dir: Path | str = PROCESSED_DIR,
) -> pl.DataFrame:
    """Подсчет пропущенных значений в каждом Parquet-файле в каталоге."""

    parquet_dir = Path(parquet_dir)
    rows = []

    for file in sorted(parquet_dir.glob("*.parquet")):
        missing_counts = (
            pl.scan_parquet(file)
            .select(pl.all().null_count())
            .collect(engine="streaming")
        )

        total_missing = sum(missing_counts.row(0))
        rows.append(
            {
                "dataset": file.name,
                "path": str(file),
                "missing_total": total_missing,
                "row_count": pl.scan_parquet(file)
                .select(pl.count())
                .collect(engine="streaming")
                .row(0)[0],
            }
        )

    return pl.DataFrame(rows).sort("missing_total")


def find_least_missing_parquet_dataset(
    parquet_dir: Path | str = PROCESSED_DIR,
) -> Path | None:
    """Возвращает путь к Parquet-файлу с наименьшим числом пропусков."""

    summary = summarize_missing_parquet_datasets(parquet_dir)
    if summary.is_empty():
        return None

    return Path(summary[0, "path"])