data_load/data_prep.py - master - NeonBite/big-data-5

data_load/data_prep.py
78 строк2 KB
NeonBite
feat: report done pushed
20 дек 2025, 17:27
20 дек 2025, 17:270df5e56
import datetime as dt
import logging
import os
import shutil

import kagglehub


DATA_DIR = os.environ.get("DATA_DIR", "data")
DATASET_NAME = os.environ.get("DATASET_NAME", "ratings_Electronics.csv")
KAGGLE_DATASET = os.environ.get("KAGGLE_DATASET", "saurav9786/amazon-product-reviews")

logger = logging.getLogger(__name__)


def setup_logging():
    log_dir = os.environ.get("LOG_DIR", ".logs")
    os.makedirs(log_dir, exist_ok=True)
    script_name = os.path.splitext(os.path.basename(__file__))[0]
    timestamp = dt.datetime.utcnow().strftime("%Y%m%d-%H%M%S")
    log_path = os.path.join(log_dir, f"{script_name}-{timestamp}.log")
    logging.basicConfig(
        level=os.environ.get("LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(name)s - %(message)s",
        handlers=[
            logging.FileHandler(log_path),
            logging.StreamHandler(),
        ],
    )
    logger.info("Logging to %s", log_path)


def normalize_dataset_name():
    base, ext = os.path.splitext(DATASET_NAME)
    alt_name = f"{base} (1){ext}"
    csv_path = os.path.join(DATA_DIR, DATASET_NAME)
    alt_path = os.path.join(DATA_DIR, alt_name)

    if not os.path.exists(csv_path) and os.path.exists(alt_path):
        logger.info("Renaming %s -> %s", alt_path, csv_path)
        shutil.move(alt_path, csv_path)


def ensure_dataset() -> str:
    os.makedirs(DATA_DIR, exist_ok=True)
    csv_path = os.path.join(DATA_DIR, DATASET_NAME)
    normalize_dataset_name()

    if not os.path.exists(csv_path) or len(os.listdir(DATA_DIR)) == 0:
        logger.info("DATA_DIR empty, downloading dataset via kagglehub")
        path = kagglehub.dataset_download(KAGGLE_DATASET)

        for name in os.listdir(path):
            src = os.path.join(path, name)
            dst = os.path.join(DATA_DIR, name)
            shutil.move(src, dst)

        logger.info("Dataset moved to %s", DATA_DIR)

    else:
        logger.info("Dataset already present in %s", DATA_DIR)

    normalize_dataset_name()

    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"Dataset not found at {csv_path}")

    return csv_path


def main():
    setup_logging()
    csv_path = ensure_dataset()
    logger.info("Dataset ready at %s", csv_path)


if __name__ == "__main__":
    main()