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