faceberg 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

faceberg/bridge.py

@@ -1,586 +0,0 @@
-"""Bridge between HuggingFace datasets and Apache Iceberg tables.
-
-This module discovers HuggingFace datasets and converts them to TableInfo objects
-that contain all the Iceberg metadata needed for table creation.
-"""
-
-import json
-import os
-import tempfile
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
-
-from datasets import (
-    Features,
-    load_dataset_builder,
-)
-from huggingface_hub import HfApi, HfFileSystem
-from pyiceberg.io.pyarrow import _pyarrow_to_schema_without_ids
-from pyiceberg.partitioning import PartitionField, PartitionSpec
-from pyiceberg.schema import Schema, assign_fresh_schema_ids
-from pyiceberg.transforms import IdentityTransform
-from pyiceberg.types import ListType, MapType, NestedField, StringType, StructType
-
-# =============================================================================
-# Bridge Output Classes
-# =============================================================================
-
-
-@dataclass
-class FileInfo:
-    """Information about a data file in Iceberg table."""
-
-    uri: str  # Full hf:// URI to the file
-    split: Optional[str] = None  # Split name (train, test, validation, etc.)
-    size_bytes: Optional[int] = None  # File size in bytes (enriched later)
-    row_count: Optional[int] = None  # Number of rows in the file (enriched later)
-
-
-@dataclass
-class TableInfo:
-    """Complete information needed to create an Iceberg table.
-
-    This class serves as the output of the bridge layer, containing all the
-    metadata needed to convert a HuggingFace dataset into an Iceberg table.
-    """
-
-    # Table identity
-    namespace: str  # Iceberg namespace (e.g., "default")
-    table_name: str  # Table name (e.g., "squad_plain_text")
-
-    # Iceberg schema and partitioning
-    schema: Schema  # Iceberg schema with field IDs
-    partition_spec: PartitionSpec  # Partition specification
-
-    # Data files
-    data_files: List[FileInfo]  # List of data files with metadata
-    data_dir: str  # Data directory path relative to repo root
-
-    # Source metadata (for traceability)
-    dataset_repo: str  # HuggingFace repo ID
-    dataset_config: str  # Dataset configuration name
-    dataset_revision: str  # Git revision/SHA of the dataset
-
-    @property
-    def identifier(self) -> str:
-        """Get table identifier in 'namespace.table_name' format."""
-        return f"{self.namespace}.{self.table_name}"
-
-    @property
-    def total_rows(self) -> int:
-        """Get total row count across all files."""
-        return sum(f.row_count for f in self.data_files if f.row_count is not None)
-
-    @property
-    def total_size(self) -> int:
-        """Get total size in bytes across all files."""
-        return sum(f.size_bytes for f in self.data_files if f.size_bytes is not None)
-
-    def get_table_properties(self) -> Dict[str, str]:
-        """Get table properties for Iceberg metadata.
-
-        Returns:
-            Dictionary of table properties including source metadata and name mapping
-        """
-        # Create schema name mapping for Parquet files without embedded field IDs
-        name_mapping = iceberg_name_mapping(self.schema)
-
-        # Use data directory from discovery
-        data_path = (
-            f"hf://datasets/{self.dataset_repo}/{self.data_dir}"
-            if self.data_dir
-            else f"hf://datasets/{self.dataset_repo}"
-        )
-
-        # TODO(kszucs): split should be configurable
-        properties = {
-            "format-version": "3",
-            "write.parquet.compression-codec": "snappy",
-            "write.py-location-provider.impl": "faceberg.catalog.HfLocationProvider",
-            "write.data.path": data_path,
-            # HuggingFace source metadata
-            "hf.dataset.repo": self.dataset_repo,
-            "hf.dataset.config": self.dataset_config,
-            "hf.dataset.revision": self.dataset_revision,
-            # Write configuration
-            "hf.write.pattern": "{split}-{uuid}-iceberg.parquet",
-            "hf.write.split": "train",
-            # Schema mapping
-            "schema.name-mapping.default": json.dumps(name_mapping),
-        }
-
-        return properties
-
-
-# =============================================================================
-# Iceberg Helpers (Schema and Metadata)
-# =============================================================================
-
-
-def iceberg_field_mapping(field: NestedField) -> Dict[str, any]:
-    """Build name mapping for a single field, recursively handling nested types.
-
-    Args:
-        field: Iceberg NestedField to create mapping for
-
-    Returns:
-        Dictionary containing field-id, names, and optionally nested fields
-    """
-    mapping = {
-        "field-id": field.field_id,
-        "names": [field.name],
-    }
-
-    # Handle nested types
-    if isinstance(field.field_type, StructType):
-        # Recursively map nested struct fields
-        nested_fields = []
-        for nested_field in field.field_type.fields:
-            nested_fields.append(iceberg_field_mapping(nested_field))
-        if nested_fields:
-            mapping["fields"] = nested_fields
-    elif isinstance(field.field_type, ListType):
-        # Create mapping for the list element
-        element_mapping = {
-            "field-id": field.field_type.element_id,
-            "names": ["element"],
-        }
-        # If element is a struct, recursively map its fields
-        if isinstance(field.field_type.element_type, StructType):
-            element_fields = []
-            for nested_field in field.field_type.element_type.fields:
-                element_fields.append(iceberg_field_mapping(nested_field))
-            if element_fields:
-                element_mapping["fields"] = element_fields
-        mapping["fields"] = [element_mapping]
-    elif isinstance(field.field_type, MapType):
-        # Create mappings for key and value
-        map_fields = []
-
-        # Map the key
-        key_mapping = {
-            "field-id": field.field_type.key_id,
-            "names": ["key"],
-        }
-        if isinstance(field.field_type.key_type, StructType):
-            key_fields = []
-            for nested_field in field.field_type.key_type.fields:
-                key_fields.append(iceberg_field_mapping(nested_field))
-            if key_fields:
-                key_mapping["fields"] = key_fields
-        map_fields.append(key_mapping)
-
-        # Map the value
-        value_mapping = {
-            "field-id": field.field_type.value_id,
-            "names": ["value"],
-        }
-        if isinstance(field.field_type.value_type, StructType):
-            value_fields = []
-            for nested_field in field.field_type.value_type.fields:
-                value_fields.append(iceberg_field_mapping(nested_field))
-            if value_fields:
-                value_mapping["fields"] = value_fields
-        map_fields.append(value_mapping)
-
-        mapping["fields"] = map_fields
-
-    return mapping
-
-
-def iceberg_name_mapping(schema: Schema) -> List[Dict[str, any]]:
-    """Build Iceberg name mapping from schema, recursively handling nested fields.
-
-    Name mapping is used to map Parquet column names to Iceberg field IDs for
-    files that don't have embedded field IDs.
-
-    Args:
-        schema: Iceberg schema to create mapping for
-
-    Returns:
-        List of field mappings with field-id, names, and nested fields
-    """
-    fields = []
-    for field in schema.fields:
-        fields.append(iceberg_field_mapping(field))
-    return fields
-
-
-def iceberg_partition_spec(schema: Schema) -> PartitionSpec:
-    """Build a partition spec that uses 'split' as a partition key.
-
-    This creates an identity partition on the split column, which means the split
-    value will be stored in metadata and used for partition pruning.
-
-    Args:
-        schema: Iceberg schema containing a 'split' field
-
-    Returns:
-        PartitionSpec with split as partition key
-
-    Raises:
-        ValueError: If schema doesn't contain a 'split' field
-    """
-    split_field = schema.find_field("split")
-    if split_field is None:
-        raise ValueError("Schema must contain a 'split' field to create split partition spec")
-
-    return PartitionSpec(
-        PartitionField(
-            source_id=split_field.field_id,
-            field_id=1000,  # Partition field IDs start at 1000
-            transform=IdentityTransform(),
-            name="split",
-        ),
-        spec_id=0,
-    )
-
-
-def iceberg_schema_from_features(features, include_split_column: bool = True) -> Schema:
-    """
-    Build an Iceberg Schema from HuggingFace dataset features using Arrow as an intermediate format.
-
-    This approach ensures globally unique field IDs across nested structures by leveraging
-    PyIceberg's built-in conversion and ID assignment logic.
-
-    Args:
-        features: HuggingFace Features object or dict of features
-        include_split_column: If True, adds a 'split' column to the schema (default: True)
-
-    Returns:
-        PyIceberg Schema object with globally unique field IDs
-    """
-    # Convert to Features if dict
-    if isinstance(features, dict):
-        features = Features(features)
-
-    # Convert: Features → Arrow Schema → Iceberg Schema (without IDs) → Assign fresh IDs
-    # This ensures globally unique field IDs across all nested structures
-    arrow_schema = features.arrow_schema
-    iceberg_schema_no_ids = _pyarrow_to_schema_without_ids(arrow_schema)
-    schema = assign_fresh_schema_ids(iceberg_schema_no_ids)
-
-    # Add split column as the first field if requested
-    if include_split_column:
-        # Create split field (will get ID 1 after reassignment)
-        # Note: Although the schema uses StringType, the actual Parquet data
-        # will use dictionary encoding (int8 indices) for compression efficiency
-        # The split column is optional since it doesn't exist in the source Parquet files,
-        # it's derived from partition metadata
-        split_field = NestedField(
-            field_id=-1,  # Temporary ID, will be reassigned
-            name="split",
-            field_type=StringType(),
-            required=False,
-        )
-        # Prepend split field to existing fields
-        new_fields = [split_field] + list(schema.fields)
-
-        # Create new schema and reassign all field IDs globally
-        # This ensures field IDs are globally unique across nested structures
-        schema_with_split = Schema(*new_fields)
-        schema = assign_fresh_schema_ids(schema_with_split)
-
-    return schema
-
-
-# =============================================================================
-# Dataset Helpers (HuggingFace)
-# =============================================================================
-
-
-def dataset_new_files(
-    repo_id: str,
-    config: str,
-    old_revision: str,
-    new_revision: str,
-    token: Optional[str] = None,
-) -> List[str]:
-    """Find new parquet files added between two revisions.
-
-    Uses HuggingFace Hub API to diff two git revisions and identify
-    new parquet files for a specific dataset configuration.
-
-    Args:
-        repo_id: HuggingFace dataset repo ID (e.g., "squad")
-        config: Dataset configuration name (e.g., "plain_text")
-        old_revision: Previous commit SHA
-        new_revision: Current commit SHA or branch (usually "main")
-        token: HuggingFace API token
-
-    Returns:
-        List of path_in_repo strings for new parquet files
-
-    Example:
-        >>> dataset_new_files(
-        ...     "squad",
-        ...     "plain_text",
-        ...     "abc123",
-        ...     "def456"
-        ... )
-        ['plain_text/train-00000-of-00001.parquet', 'plain_text/validation-00000-of-00001.parquet']
-    """
-    api = HfApi(token=token)
-
-    # Get all files at old revision
-    old_files = set(
-        api.list_repo_files(
-            repo_id=repo_id,
-            repo_type="dataset",
-            revision=old_revision,
-        )
-    )
-
-    # Get all files at new revision
-    new_files = set(
-        api.list_repo_files(
-            repo_id=repo_id,
-            repo_type="dataset",
-            revision=new_revision,
-        )
-    )
-
-    # Find added files (set difference)
-    added_files = new_files - old_files
-
-    # Filter for parquet files in this config
-    config_prefix = f"{config}/"
-    relative_paths = sorted(
-        f for f in added_files if f.endswith(".parquet") and f.startswith(config_prefix)
-    )
-
-    return relative_paths
-
-
-def dataset_builder_safe(
-    repo_id: str,
-    config: str,
-    token: Optional[str] = None,
-):
-    """Load dataset builder while avoiding picking up local files.
-
-    Changes to a temporary directory before loading to ensure the datasets
-    library doesn't pick up local files in the current directory.
-
-    Args:
-        repo_id: HuggingFace dataset repository ID
-        config_name: Optional configuration name
-        token: Optional HuggingFace API token
-
-    Returns:
-        Dataset builder object
-
-    Raises:
-        Exception: If loading fails
-    """
-    original_cwd = os.getcwd()
-
-    try:
-        # Change to a temporary directory to avoid dataset library picking up local files
-        with tempfile.TemporaryDirectory() as tmpdir:
-            os.chdir(tmpdir)
-            return load_dataset_builder(repo_id, config, token=token)
-    finally:
-        # Always restore the original directory
-        os.chdir(original_cwd)
-
-
-def dataset_data_files(
-    data_files: Dict[str, List[str]], filter_paths: Optional[List[str]] = None
-) -> Tuple[Dict[str, List[str]], str]:
-    """Filter data files and extract data directory.
-
-    Optionally filters data files by path_in_repo and extracts the common directory path.
-
-    Args:
-        data_files: Dictionary mapping splits to lists of file URIs (with revision)
-        filter_paths: Optional list of path_in_repo strings to filter by
-
-    Returns:
-        Tuple of (filtered_data_files, data_dir)
-
-    Example:
-        >>> dataset_data_files(
-        ...     {"train": ["hf://datasets/repo@rev/plain_text/train-00000.parquet"]},
-        ...     filter_paths=["plain_text/train-00000.parquet"]
-        ... )
-        ({'train': ['hf://datasets/repo@rev/plain_text/train-00000.parquet']}, 'plain_text')
-    """
-    fs = HfFileSystem()
-
-    # Convert filter_paths to set for fast lookup
-    filter_set = set(filter_paths) if filter_paths else None
-
-    # Filter data files if filter_paths provided
-    filtered_data_files = {}
-    all_files = []
-
-    for split, file_list in data_files.items():
-        filtered_files = []
-        for file_uri in file_list:
-            resolved = fs.resolve_path(file_uri)
-            path_in_repo = resolved.path_in_repo
-
-            # Include file if no filter or if in filter set
-            if filter_set is None or path_in_repo in filter_set:
-                filtered_files.append(file_uri)
-                all_files.append(path_in_repo)
-
-        if filtered_files:
-            filtered_data_files[split] = filtered_files
-
-    if not all_files:
-        raise ValueError("No data files found to determine data directory")
-
-    try:
-        # Extract directory from each file path, then find common directory
-        # This ensures we get a directory path, not a file path (which would happen
-        # with os.path.commonpath when there's only one file)
-        directories = [os.path.dirname(f) for f in all_files]
-        data_dir = os.path.commonpath(directories)
-    except ValueError as e:
-        raise ValueError(
-            f"Unable to determine common data directory from files: {all_files}"
-        ) from e
-
-    return filtered_data_files, data_dir
-
-
-# =============================================================================
-# Dataset Discovery and Bridging
-# =============================================================================
-
-
-@dataclass
-class DatasetInfo:
-    """Information about a HuggingFace dataset.
-
-    This class discovers and represents the structure of a HuggingFace dataset,
-    including its configuration, splits, and Parquet files. It serves as the
-    discovery layer that gathers all necessary information before conversion
-    to Iceberg format.
-    """
-
-    repo_id: str
-    config: str
-    splits: List[str]
-    data_files: Dict[str, List[str]]  # split -> list of fully qualified URIs (with revision)
-    data_dir: str
-    features: Features  # HuggingFace dataset features
-    revision: str  # Git revision/SHA of the dataset
-
-    @classmethod
-    def discover(
-        cls,
-        repo_id: str,
-        config: str,
-        token: Optional[str] = None,
-        since_revision: Optional[str] = None,
-    ) -> "DatasetInfo":
-        """Discover Parquet files and structure in a HuggingFace dataset.
-
-        Discovery process:
-        1. Validate config exists in dataset
-        2. Load dataset builder to get metadata
-        3. Extract splits from builder
-        4. Get data files (fully qualified URIs with revision)
-           - If since_revision is provided, only get files added since that revision
-           - Otherwise, get all files from builder
-        5. Get dataset revision (SHA) from Hub
-        6. Extract data directory from config or URIs
-        7. Return DatasetInfo with all metadata
-
-        Args:
-            repo_id: HuggingFace dataset repository ID (e.g., "kszucs/dataset1")
-            config: Configuration name to discover
-            token: HuggingFace API token (uses HF_TOKEN env var if not provided)
-            since_revision: Optional revision SHA to get only files added since that revision
-
-        Returns:
-            DatasetInfo with discovered structure
-
-        Raises:
-            ValueError: If dataset not found or config doesn't exist
-        """
-        try:
-            builder = dataset_builder_safe(repo_id, config=config, token=token)
-        except Exception as e:
-            raise ValueError(
-                f"Dataset {repo_id} config {config} not found or not accessible: {e}"
-            ) from e
-
-        revision = builder.hash
-        features = builder.info.features
-
-        # Get filter paths if since_revision is provided
-        filter_paths = None
-        if since_revision:
-            filter_paths = dataset_new_files(
-                repo_id=repo_id,
-                config=config,
-                old_revision=since_revision,
-                new_revision=revision,
-                token=token,
-            )
-
-        # Filter data files and extract data directory
-        data_files, data_dir = dataset_data_files(
-            builder.config.data_files, filter_paths=filter_paths
-        )
-
-        splits = list(data_files.keys())
-
-        if not data_files:
-            raise ValueError("No Parquet files found in dataset configuration")
-
-        return cls(
-            repo_id=repo_id,
-            config=config,
-            splits=splits,
-            data_files=data_files,  # Store fully qualified URIs
-            data_dir=data_dir,
-            features=features,
-            revision=revision,
-        )
-
-    def to_table_info(
-        self,
-        namespace: str,
-        table_name: str,
-    ) -> TableInfo:
-        """Convert DatasetInfo to TableInfo.
-
-        This method creates table metadata for the HuggingFace dataset config
-        with an explicit table name, supporting the namespace-based configuration.
-
-        Args:
-            namespace: Iceberg namespace for the table
-            table_name: Explicit table name (no auto-generation)
-
-        Returns:
-            TableInfo object
-        """
-        # Build Iceberg schema with split column
-        schema = iceberg_schema_from_features(self.features, include_split_column=True)
-
-        # Build partition spec (partitioned by split)
-        partition_spec = iceberg_partition_spec(schema)
-
-        # Collect file information with fully qualified URIs
-        files = []
-        for split, file_uris in self.data_files.items():
-            for uri in file_uris:
-                files.append(FileInfo(uri=uri, split=split))
-
-        # Create TableInfo with explicit naming
-        return TableInfo(
-            namespace=namespace,
-            table_name=table_name,  # Direct from config, no auto-generation
-            schema=schema,
-            partition_spec=partition_spec,
-            data_files=files,
-            data_dir=self.data_dir,
-            dataset_repo=self.repo_id,
-            dataset_config=self.config,
-            dataset_revision=self.revision,
-        )