dataset-uploader/scripts/dataset_validator.py

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

from __future__ import annotations

import csv
import logging
import os
from collections.abc import Callable, Iterator
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

from scripts.download_and_decompress import decompress, download

logger = logging.getLogger(__name__)
logging.basicConfig(filename="validator.log", level=logging.DEBUG)
logger.setLevel(logging.DEBUG)


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 _generate_fields_from_source(
    dataset_id: str,
    base_dir: Path,
) -> Iterator[str | None]:
    """Generate all field values from source dataset files.

    Args:
        dataset_id: Dataset identifier.
        base_dir: Base directory containing downloads.

    Yields:
        Raw string values from CSV/TSV and RDF files (unfiltered).

    Raises:
        RuntimeError: If download or decompression fails.
    """
    download_dir = Path(base_dir) / "downloads" / dataset_id

    if not download_dir.exists():
        print(f"No downloads found for {dataset_id}. Downloading.")
        (successful, _try_again, rdf_file) = download(
            dataset_id=dataset_id, base_dir=base_dir
        )

        if not successful:
            raise RuntimeError(f"Download failed for {dataset_id}")

        if rdf_file is None:
            raise RuntimeError(f"rdf_file is None for {dataset_id}")
        decompress_result = decompress(rdf_file)
        if not decompress_result:
            raise RuntimeError(f"Decompression failed for {dataset_id}")

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

    for root, _dirs, files in os.walk(download_dir):
        for file in files:
            file_path = Path(root) / file

            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:
                            logger.debug(f"Yielding {field} from CSV/TSV file.")
                            yield field
                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:
                    logger.debug(f"Yielding {subject} from RDF file.")
                    yield str(subject)
                    logger.debug(f"Yielding {predicate} from RDF file.")
                    yield str(predicate)
                    logger.debug(f"Yielding {obj} from RDF file.")
                    yield str(obj)


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()

    try:
        for field in _generate_fields_from_source(dataset_id, base_dir):
            _add_string(store, field, filter_fn)
    except RuntimeError as e:
        print(f"Error generating fields: {e}. Returning empty store.")
        return store

    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:
    """
    Collects strings from a parquet file and returns them as a BloomFilter.

    This function reads the specified parquet file, processes its required
    columns as strings, and applies an optional filter function to determine
    whether each string value should be included in the resulting BloomFilter.
    If no filter function is provided, a default one is used.

    Parameters:
        parquet_path (Path): The path to the parquet file to be read.
        filter_fn (Callable[[str], bool] | None): An optional callable to filter
            string values. If not provided, a default filter function is used.

    Returns:
        BloomFilter: A BloomFilter containing the processed and filtered string values.
    """
    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:
            logger.debug(f"Adding {value} to parquet store.")
            _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)


def strings_in_parquet_not_source(
    parquet_paths: list[Path],
    source_strings: BloomFilter,
    filter_fn: Callable[[str], bool] | None = None,
    max_count: int = 20,
) -> list[str]:
    """
    Determines strings in a given parquet file that are not part of a specified
    source strings set.

    The function processes a parquet file, applying a filtering function to identify
    strings that are not in the provided `source_strings` BloomFilter. It iterates
    through specified columns in the parquet file, converts their values to strings,
    and checks them against the filter function and the `source_strings`. Only unique
    strings that match the filtering criteria and are not in the `source_strings` are
    returned up to the specified `max_count`.

    Parameters:
    parquet_paths : list[Path]
        The file paths of the parquet file to be processed.
    source_strings : BloomFilter
        A BloomFilter containing strings that should be excluded from the results.
    filter_fn : Callable[[str], bool] | None, optional
        A callable filter function that takes a string as input and returns a boolean.
        If not provided, a default filtering function is used.
    max_count : int
        The maximum number of strings to return. Defaults to 20.

    Returns:
    list[str]
        A list of strings from the parquet file that are not in the provided source
        strings and satisfy the filter function. The length of the list is limited by
        `max_count`.
    """
    result = []
    count = 0
    filter_fn = filter_fn or filter_string

    for parquet_path in parquet_paths:
        dataframe = read_parquet(parquet_path, columns=list(REQUIRED_COLUMNS))
        for column in REQUIRED_COLUMNS:
            series = dataframe[column].dropna().astype(str)
            for value in series:
                if filter_fn(value) and value not in source_strings:
                    result.append(value)
                    count += 1
                    if count >= max_count:
                        return result
    return result


def strings_in_source_not_parquet(
    dataset_id: str,
    base_dir: Path,
    parquet_strings: BloomFilter,
    filter_fn: Callable[[str], bool] | None = None,
    max_count: int = 20,
) -> list[str]:
    """
    Find strings in the source data that are not in the provided parquet-based
    BloomFilter.

    This function identifies strings from a source dataset, checks them against a
    BloomFilter built on parquet data, and collects those that are not present.
    Only strings matching a specific filter function are considered, and the
    function can optionally limit the number of returned strings.

    Parameters:
    dataset_id: str
        Identifier of the dataset to search for source strings.
    base_dir: Path
        Base directory where the source datasets are located.
    parquet_strings: BloomFilter
        BloomFilter object containing strings obtained from parquet files to check
        against.
    filter_fn: Callable[[str], bool] | None
        Function to determine which strings from the source should be considered.
        If None, a default filtering function (filter_string) is used.
    max_count: int
        Maximum number of strings to find in the source dataset before stopping.
        Defaults to 20.

    Returns:
    list[str]
        A list of strings found in the source data that are not present in the
        parquet BloomFilter. The number of strings is limited by the `max_count`
        parameter.
    """
    result = []
    count = 0
    filter_fn = filter_fn or filter_string

    try:
        for field in _generate_fields_from_source(dataset_id, base_dir):
            if field is not None and filter_fn(field) and field not in parquet_strings:
                result.append(field)
                count += 1
                if count >= max_count:
                    return result
    except RuntimeError as e:
        print(f"Error generating fields: {e}. Returning empty list.")
        return []

    return result