dataset-uploader/scripts/dataset_validator.py

"""Dataset validation utilities for HuggingFace uploads."""

from __future__ import annotations

import csv
from collections.abc import Callable
from pathlib import Path

from bloom_filter2 import BloomFilter
from dataset_registry import get_dataset_config
from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.errors import HfHubHTTPError
from pandas import read_parquet
from rdflib import Graph


def resolve_organization(organization: str | None = None) -> str:
    """Resolve the HuggingFace organization name."""
    if organization:
        return organization

    config = get_dataset_config()
    return config.get("organization", "CleverThis")


def build_repo_id(dataset_id: str, organization: str | None = None) -> str:
    """Build the full HuggingFace repo ID for a dataset."""
    org_name = resolve_organization(organization)
    return f"{org_name}/{dataset_id}"


def dataset_exists_on_hf(
    dataset_id: str,
    organization: str | None = None,
    api: HfApi | None = None,
) -> bool:
    """Check whether a dataset exists on HuggingFace Hub.

    Args:
        dataset_id: Dataset identifier (without organization prefix).
        organization: HuggingFace organization name.
        api: Optional HfApi instance for dependency injection.

    Returns:
        True if dataset exists, False if not found.

    Raises:
        HfHubHTTPError: For non-404 errors.
    """
    repo_id = build_repo_id(dataset_id, organization)
    client = api or HfApi()

    try:
        client.dataset_info(repo_id=repo_id)
    except HfHubHTTPError as exc:
        status_code = exc.response.status_code if exc.response else None
        if status_code is not None and status_code >= 400:
            return False
        raise

    return True


def dataset_can_download_from_hf(
    dataset_id: str,
    organization: str | None = None,
    download_func: Callable[..., str] | None = None,
    cache_dir: Path | str | None = None,
) -> str | None:
    """Check whether a dataset can be downloaded from HuggingFace Hub.

    Args:
        dataset_id: Dataset identifier (without organization prefix).
        organization: HuggingFace organization name.
        download_func: Optional download function for dependency injection.
        cache_dir: Optional cache directory for HuggingFace downloads.

    Returns:
        Download path if download succeeds, None if not found or download fails.

    Raises:
        HfHubHTTPError: For non-404 errors.
    """
    repo_id = build_repo_id(dataset_id, organization)
    downloader = download_func or snapshot_download

    try:
        download_path = downloader(
            repo_id=repo_id,
            repo_type="dataset",
            local_dir=cache_dir,
        )
    except HfHubHTTPError as exc:
        status_code = exc.response.status_code if exc.response else None
        if status_code == 404:
            return None
        raise
    except Exception:
        return None

    return download_path


def dataset_is_parquet_file(
    file_path: Path | str,
    read_parquet_func: Callable[[Path | str], object] | None = None,
) -> bool:
    """Check whether a file can be read by pandas.read_parquet().

    Args:
        file_path: Path to the file to validate.
        read_parquet_func: Optional read_parquet function for dependency injection.

    Returns:
        True if the file can be read as parquet, False otherwise.
    """
    reader = read_parquet_func or read_parquet

    try:
        reader(file_path)
    except Exception:
        return False

    return True


def _new_string_store() -> BloomFilter:
    return BloomFilter(max_elements=10000000, error_rate=0.001)


def _add_string(
    store: BloomFilter,
    value: str | None,
    filter_fn: Callable[[str], bool],
) -> None:
    if value is None:
        return
    if filter_fn(value):
        store.add(value)


def filter_string(value: str | None) -> bool:
    if value is None:
        return False
    return value.strip() != ""


def collect_source_strings(
    dataset_id: str,
    base_dir: Path,
    filter_fn: Callable[[str], bool] | None = None,
) -> BloomFilter:
    filter_fn = filter_fn or filter_string
    store = _new_string_store()
    download_dir = Path(base_dir) / "downloads" / dataset_id

    if not download_dir.exists():
        return store

    compressed_suffixes = {".gz", ".bz2", ".zip", ".xz", ".tgz", ".tar"}

    for file_path in download_dir.rglob("*"):
        if not file_path.is_file():
            continue

        if any(suffix in compressed_suffixes for suffix in file_path.suffixes):
            continue

        suffix = file_path.suffix.lower()
        if suffix in {".csv", ".tsv"}:
            delimiter = "\t" if suffix == ".tsv" else ","
            with open(file_path, newline="", encoding="utf-8") as file_obj:
                reader = csv.reader(file_obj, delimiter=delimiter)
                for row in reader:
                    for field in row:
                        _add_string(store, field, filter_fn)
            continue

        if suffix in {".nt", ".ntriples", ".ttl", ".turtle", ".rdf", ".xml"}:
            if suffix in {".nt", ".ntriples"}:
                rdf_format = "nt"
            elif suffix in {".ttl", ".turtle"}:
                rdf_format = "turtle"
            else:
                rdf_format = "xml"

            graph = Graph()
            graph.parse(str(file_path), format=rdf_format)
            for subject, predicate, obj in graph:
                _add_string(store, str(subject), filter_fn)
                _add_string(store, str(predicate), filter_fn)
                _add_string(store, str(obj), filter_fn)

    return store

REQUIRED_COLUMNS = (
    "subject",
    "predicate",
    "object",
    "object_type",
    "object_datatype",
    "object_language",
)

def collect_parquet_strings(
    parquet_path: Path,
    filter_fn: Callable[[str], bool] | None = None,
) -> BloomFilter:
    filter_fn = filter_fn or filter_string
    store = _new_string_store()

    dataframe = read_parquet(parquet_path, columns=list(REQUIRED_COLUMNS))
    for column in REQUIRED_COLUMNS:
        series = dataframe[column].dropna().astype(str)
        for value in series:
            _add_string(store, value, filter_fn)

    return store


def dataset_has_required_columns(file_path: Path | str) -> bool:
    """Check whether a parquet file has the required RDF columns."""
    try:
        dataframe = read_parquet(file_path)
    except Exception:
        return False

    columns = set(dataframe.columns)
    return all(column in columns for column in REQUIRED_COLUMNS)