earthcatalog 0.2.0__py3-none-any.whl

earthcatalog/pipeline.py

@@ -0,0 +1,606 @@
+# pipeline.py
+"""Pipeline components for batch processing and shard consolidation.
+
+This module provides reusable components for the STAC ingestion pipeline,
+including configuration dataclasses and helper utilities for batch processing
+and shard consolidation operations.
+
+Components:
+    BatchConfig: Configuration for URL batch processing
+    ConsolidationConfig: Configuration for shard consolidation
+    ShardInfo: Metadata about a written shard
+    PartitionResult: Result of consolidating a partition
+    BatchResult: Result of processing a URL batch
+
+These components are designed to work with the main STACIngestionPipeline
+while providing clear, typed interfaces for configuration and results.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+import geopandas as gpd
+import pandas as pd
+
+if TYPE_CHECKING:
+    from .storage_backends import StorageBackend
+
+# Type alias for duplicate handling
+DropKeep = Literal["first", "last"]
+CompressionType = Literal["snappy", "gzip", "brotli"]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BatchConfig:
+    """Configuration for URL batch processing.
+
+    Controls how URLs are chunked and processed across workers,
+    including async HTTP settings and memory management.
+
+    Attributes:
+        batch_size: Number of URLs to process in each async batch.
+        items_per_shard: Target number of items per shard file.
+        enable_concurrent_http: Whether to use async HTTP processing.
+        concurrent_requests: Number of concurrent HTTP requests.
+        connection_pool_size: Size of the HTTP connection pool.
+        request_timeout: Timeout for each HTTP request in seconds.
+        retry_attempts: Number of retry attempts for failed requests.
+        retry_delay: Base delay between retries in seconds.
+
+    Example:
+        >>> config = BatchConfig(
+        ...     batch_size=1000,
+        ...     concurrent_requests=50,
+        ...     enable_concurrent_http=True
+        ... )
+        >>> print(f"Processing {config.batch_size} URLs per batch")
+    """
+
+    batch_size: int = 1000
+    items_per_shard: int = 10000
+    enable_concurrent_http: bool = True
+    concurrent_requests: int = 50
+    connection_pool_size: int = 100
+    request_timeout: int = 30
+    retry_attempts: int = 3
+    retry_delay: float = 1.0
+
+    def __post_init__(self) -> None:
+        """Validate configuration after initialization."""
+        self._validate()
+
+    def _validate(self) -> None:
+        """Validate configuration values."""
+        if self.batch_size <= 0:
+            raise ValueError("batch_size must be positive")
+        if self.items_per_shard <= 0:
+            raise ValueError("items_per_shard must be positive")
+        if self.concurrent_requests <= 0:
+            raise ValueError("concurrent_requests must be positive")
+        if self.connection_pool_size <= 0:
+            raise ValueError("connection_pool_size must be positive")
+        if self.request_timeout <= 0:
+            raise ValueError("request_timeout must be positive")
+        if self.retry_attempts < 0:
+            raise ValueError("retry_attempts must be non-negative")
+        if self.retry_delay < 0:
+            raise ValueError("retry_delay must be non-negative")
+
+    def __repr__(self) -> str:
+        """Return detailed string representation."""
+        return (
+            f"BatchConfig(batch_size={self.batch_size}, "
+            f"items_per_shard={self.items_per_shard}, "
+            f"enable_concurrent_http={self.enable_concurrent_http}, "
+            f"concurrent_requests={self.concurrent_requests})"
+        )
+
+    def __bool__(self) -> bool:
+        """Return True if configuration is valid."""
+        try:
+            self._validate()
+            return True
+        except ValueError:
+            return False
+
+
+@dataclass
+class ConsolidationConfig:
+    """Configuration for shard consolidation operations.
+
+    Controls memory management, merge strategies, and output settings
+    for consolidating worker shards into final partitioned catalogs.
+
+    Attributes:
+        strategy: Consolidation strategy - "efficient" or "legacy".
+        max_memory_per_partition_mb: Memory limit per partition in MB.
+        enable_streaming_merge: Use streaming for large file merges.
+        s3_multipart_threshold_mb: Threshold for S3 multipart uploads.
+        temp_dir_location: Location for temporary staging files.
+        sort_key: Column to sort consolidated data by.
+        sort_ascending: Sort order (True for ascending).
+        deduplicate_key: Column to use for deduplication.
+        keep_duplicates: Which duplicate to keep ("first" or "last").
+
+    Example:
+        >>> config = ConsolidationConfig(
+        ...     strategy="efficient",
+        ...     max_memory_per_partition_mb=1024,
+        ...     enable_streaming_merge=True
+        ... )
+        >>> print(f"Using {config.strategy} consolidation")
+    """
+
+    strategy: str = "efficient"
+    max_memory_per_partition_mb: int = 1024
+    enable_streaming_merge: bool = True
+    s3_multipart_threshold_mb: int = 100
+    temp_dir_location: str = field(default_factory=tempfile.gettempdir)
+    sort_key: str = "datetime"
+    sort_ascending: bool = True
+    deduplicate_key: str = "id"
+    keep_duplicates: str = "last"
+
+    def __post_init__(self) -> None:
+        """Validate configuration after initialization."""
+        self._validate()
+
+    def _validate(self) -> None:
+        """Validate configuration values."""
+        if self.strategy not in ("efficient", "legacy"):
+            raise ValueError("strategy must be 'efficient' or 'legacy'")
+        if self.max_memory_per_partition_mb <= 0:
+            raise ValueError("max_memory_per_partition_mb must be positive")
+        if self.s3_multipart_threshold_mb <= 0:
+            raise ValueError("s3_multipart_threshold_mb must be positive")
+        if self.keep_duplicates not in ("first", "last"):
+            raise ValueError("keep_duplicates must be 'first' or 'last'")
+
+    def __repr__(self) -> str:
+        """Return detailed string representation."""
+        return (
+            f"ConsolidationConfig(strategy='{self.strategy}', "
+            f"max_memory_mb={self.max_memory_per_partition_mb}, "
+            f"streaming={self.enable_streaming_merge})"
+        )
+
+    def __bool__(self) -> bool:
+        """Return True if configuration is valid."""
+        try:
+            self._validate()
+            return True
+        except ValueError:
+            return False
+
+
+@dataclass
+class ShardInfo:
+    """Metadata about a written shard file.
+
+    Captures information about a shard written during batch processing,
+    including its location, size, and partition assignment.
+
+    Attributes:
+        shard_path: Full path to the shard file.
+        partition_key: Partition key this shard belongs to.
+        item_count: Number of items in the shard.
+        worker_id: ID of the worker that created this shard.
+        shard_id: Sequential ID within the worker's shards.
+
+    Example:
+        >>> shard = ShardInfo(
+        ...     shard_path="/scratch/shards/h3_82/worker-0.parquet",
+        ...     partition_key="dataset/partition=h3/level=2/82/year=2023/month=01",
+        ...     item_count=500,
+        ...     worker_id="worker-0-abc123"
+        ... )
+        >>> print(f"Shard has {shard.item_count} items")
+    """
+
+    shard_path: str
+    partition_key: str
+    item_count: int
+    worker_id: str
+    shard_id: int = 0
+
+    def __repr__(self) -> str:
+        """Return detailed string representation."""
+        return f"ShardInfo(path='{self.shard_path}', partition='{self.partition_key}', items={self.item_count})"
+
+    def __bool__(self) -> bool:
+        """Return True if shard has items."""
+        return self.item_count > 0
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "shard_path": self.shard_path,
+            "partition_key": self.partition_key,
+            "item_count": self.item_count,
+            "worker_id": self.worker_id,
+            "shard_id": self.shard_id,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ShardInfo:
+        """Create ShardInfo from dictionary."""
+        return cls(
+            shard_path=data["shard_path"],
+            partition_key=data.get("partition_key", ""),
+            item_count=data["item_count"],
+            worker_id=data["worker_id"],
+            shard_id=data.get("shard_id", 0),
+        )
+
+
+@dataclass
+class PartitionResult:
+    """Result of consolidating a single partition.
+
+    Captures the outcome of merging shards into a final partition file,
+    including counts of existing, new, and deduplicated items.
+
+    Attributes:
+        partition_key: The partition that was consolidated.
+        item_count: Total items in the final partition.
+        existing_count: Items that existed before consolidation.
+        new_count: New items added during consolidation.
+        duplicates_removed: Number of duplicate items removed.
+        final_path: Path to the consolidated partition file.
+        success: Whether consolidation succeeded.
+        error: Error message if consolidation failed.
+
+    Example:
+        >>> result = PartitionResult(
+        ...     partition_key="dataset/partition=h3/level=2/82/year=2023/month=01",
+        ...     item_count=1500,
+        ...     existing_count=1000,
+        ...     new_count=550,
+        ...     duplicates_removed=50
+        ... )
+        >>> print(f"Partition has {result.item_count} total items")
+    """
+
+    partition_key: str
+    item_count: int = 0
+    existing_count: int = 0
+    new_count: int = 0
+    duplicates_removed: int = 0
+    final_path: str = ""
+    success: bool = True
+    error: str = ""
+
+    def __repr__(self) -> str:
+        """Return detailed string representation."""
+        status = "OK" if self.success else f"FAILED: {self.error}"
+        return (
+            f"PartitionResult(partition='{self.partition_key}', "
+            f"items={self.item_count}, new={self.new_count}, "
+            f"deduped={self.duplicates_removed}, status={status})"
+        )
+
+    def __bool__(self) -> bool:
+        """Return True if consolidation was successful with items."""
+        return self.success and self.item_count > 0
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "partition": self.partition_key,
+            "item_count": self.item_count,
+            "existing_count": self.existing_count,
+            "new_count": self.new_count,
+            "duplicates_removed": self.duplicates_removed,
+            "final_path": self.final_path,
+            "success": self.success,
+            "error": self.error,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> PartitionResult:
+        """Create PartitionResult from dictionary."""
+        return cls(
+            partition_key=data.get("partition", ""),
+            item_count=data.get("item_count", 0),
+            existing_count=data.get("existing_count", 0),
+            new_count=data.get("new_count", 0),
+            duplicates_removed=data.get("duplicates_removed", 0),
+            final_path=data.get("final_path", ""),
+            success=data.get("success", True),
+            error=data.get("error", ""),
+        )
+
+    @classmethod
+    def empty(cls, partition_key: str) -> PartitionResult:
+        """Create an empty result for a partition with no items."""
+        return cls(partition_key=partition_key, item_count=0)
+
+    @classmethod
+    def failed(cls, partition_key: str, error: str) -> PartitionResult:
+        """Create a failed result for a partition."""
+        return cls(partition_key=partition_key, success=False, error=error)
+
+
+@dataclass
+class BatchResult:
+    """Result of processing a batch of URLs.
+
+    Captures the outcome of processing a URL batch, including
+    the shards created and statistics collected.
+
+    Attributes:
+        worker_id: ID of the worker that processed this batch.
+        shards: List of ShardInfo for shards created.
+        urls_processed: Total URLs attempted.
+        urls_succeeded: URLs successfully processed.
+        urls_failed: URLs that failed processing.
+        stats: Statistics collected during processing.
+
+    Example:
+        >>> result = BatchResult(
+        ...     worker_id="worker-0",
+        ...     shards=[shard1, shard2],
+        ...     urls_processed=1000,
+        ...     urls_succeeded=995,
+        ...     urls_failed=5
+        ... )
+        >>> print(f"Success rate: {result.success_rate:.1%}")
+    """
+
+    worker_id: str
+    shards: list[ShardInfo] = field(default_factory=list)
+    urls_processed: int = 0
+    urls_succeeded: int = 0
+    urls_failed: int = 0
+    stats: Any = None  # IngestionStatistics when available
+
+    @property
+    def success_rate(self) -> float:
+        """Calculate the success rate for this batch."""
+        if self.urls_processed == 0:
+            return 0.0
+        return self.urls_succeeded / self.urls_processed
+
+    @property
+    def total_items(self) -> int:
+        """Total items across all shards."""
+        return sum(shard.item_count for shard in self.shards)
+
+    def __repr__(self) -> str:
+        """Return detailed string representation."""
+        return (
+            f"BatchResult(worker='{self.worker_id}', "
+            f"shards={len(self.shards)}, items={self.total_items}, "
+            f"success_rate={self.success_rate:.1%})"
+        )
+
+    def __bool__(self) -> bool:
+        """Return True if batch produced any items."""
+        return self.total_items > 0
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "worker_id": self.worker_id,
+            "shards": [s.to_dict() for s in self.shards],
+            "urls_processed": self.urls_processed,
+            "urls_succeeded": self.urls_succeeded,
+            "urls_failed": self.urls_failed,
+            "total_items": self.total_items,
+            "success_rate": self.success_rate,
+        }
+
+
+def merge_geodataframes(
+    dataframes: list[gpd.GeoDataFrame],
+    deduplicate_key: str = "id",
+    keep: DropKeep = "last",
+    sort_key: str | None = None,
+    sort_ascending: bool = True,
+) -> gpd.GeoDataFrame:
+    """Merge multiple GeoDataFrames with deduplication and sorting.
+
+    Utility function for consolidation operations that combines multiple
+    GeoDataFrames, removes duplicates, and optionally sorts the result.
+
+    Args:
+        dataframes: List of GeoDataFrames to merge.
+        deduplicate_key: Column to use for deduplication.
+        keep: Which duplicate to keep - "first" or "last".
+        sort_key: Optional column to sort by after merging.
+        sort_ascending: Sort order if sort_key is specified.
+
+    Returns:
+        Merged GeoDataFrame with duplicates removed.
+
+    Example:
+        >>> merged = merge_geodataframes(
+        ...     [existing_gdf, new_gdf],
+        ...     deduplicate_key="id",
+        ...     keep="last",
+        ...     sort_key="datetime"
+        ... )
+    """
+    if not dataframes:
+        return gpd.GeoDataFrame()
+
+    if len(dataframes) == 1:
+        merged = dataframes[0].copy()
+    else:
+        merged = pd.concat(dataframes, ignore_index=True)
+
+    # Deduplicate
+    if deduplicate_key in merged.columns:
+        original_count = len(merged)
+        merged = merged.drop_duplicates(subset=[deduplicate_key], keep=keep)
+        duplicates_removed = original_count - len(merged)
+        if duplicates_removed > 0:
+            logger.debug(f"Removed {duplicates_removed} duplicate items")
+
+    # Sort if requested
+    if sort_key and sort_key in merged.columns:
+        merged = merged.sort_values(sort_key, ascending=sort_ascending)
+
+    return gpd.GeoDataFrame(merged)
+
+
+def read_parquet_from_storage(
+    storage: StorageBackend,
+    path: str,
+) -> gpd.GeoDataFrame:
+    """Read a GeoParquet file from storage backend.
+
+    Handles the complexity of reading binary data from various storage
+    backends and converting to GeoDataFrame.
+
+    Args:
+        storage: Storage backend to read from.
+        path: Path to the parquet file.
+
+    Returns:
+        GeoDataFrame with the file contents.
+
+    Raises:
+        IOError: If file cannot be read.
+
+    Example:
+        >>> from earthcatalog.storage_backends import get_storage_backend
+        >>> storage = get_storage_backend("s3://bucket/catalog")
+        >>> gdf = read_parquet_from_storage(storage, "s3://bucket/catalog/data.parquet")
+    """
+    import pyarrow.parquet as pq
+
+    with storage.open(path, "rb") as f:
+        binary_data = f.read()
+        table = pq.read_table(io.BytesIO(binary_data))
+        df = table.to_pandas()
+        return gpd.GeoDataFrame(df)
+
+
+def write_parquet_to_storage(
+    gdf: gpd.GeoDataFrame,
+    storage: StorageBackend,
+    path: str,
+    compression: CompressionType = "snappy",
+) -> None:
+    """Write a GeoDataFrame to storage backend as GeoParquet.
+
+    Handles the complexity of writing binary data to various storage
+    backends with proper compression settings.
+
+    Args:
+        gdf: GeoDataFrame to write.
+        storage: Storage backend to write to.
+        path: Destination path for the parquet file.
+        compression: Compression codec to use (default: snappy).
+
+    Raises:
+        IOError: If file cannot be written.
+
+    Example:
+        >>> from earthcatalog.storage_backends import get_storage_backend
+        >>> storage = get_storage_backend("s3://bucket/catalog")
+        >>> write_parquet_to_storage(gdf, storage, "s3://bucket/catalog/data.parquet")
+    """
+    storage.makedirs(Path(path).parent)
+    with storage.open(path, "wb") as f:
+        gdf.to_parquet(f, index=False, compression=compression)
+
+
+def group_shards_by_partition(
+    shards: list[ShardInfo | dict[str, Any]],
+) -> dict[str, list[str]]:
+    """Group shard paths by their partition key.
+
+    Utility function that organizes shards by their target partition
+    for efficient consolidation processing.
+
+    Args:
+        shards: List of ShardInfo objects or dictionaries with shard metadata.
+
+    Returns:
+        Dictionary mapping partition keys to lists of shard paths.
+
+    Example:
+        >>> shards = [
+        ...     ShardInfo(shard_path="/a.parquet", partition_key="p1", item_count=10, worker_id="w1"),
+        ...     ShardInfo(shard_path="/b.parquet", partition_key="p1", item_count=20, worker_id="w2"),
+        ...     ShardInfo(shard_path="/c.parquet", partition_key="p2", item_count=15, worker_id="w1"),
+        ... ]
+        >>> groups = group_shards_by_partition(shards)
+        >>> print(groups)
+        {'p1': ['/a.parquet', '/b.parquet'], 'p2': ['/c.parquet']}
+    """
+    partition_shards: dict[str, list[str]] = {}
+
+    for shard in shards:
+        if isinstance(shard, ShardInfo):
+            partition_key = shard.partition_key
+            shard_path = shard.shard_path
+        else:
+            partition_key = shard.get("partition_key", "")
+            shard_path = shard.get("shard_path", "")
+
+        if not partition_key:
+            continue
+
+        if partition_key not in partition_shards:
+            partition_shards[partition_key] = []
+        partition_shards[partition_key].append(shard_path)
+
+    return partition_shards
+
+
+def chunk_urls(urls: list[str], num_chunks: int) -> list[list[str]]:
+    """Split URLs into approximately equal chunks for parallel processing.
+
+    Args:
+        urls: List of URLs to chunk.
+        num_chunks: Number of chunks to create.
+
+    Returns:
+        List of URL lists, one per chunk.
+
+    Example:
+        >>> urls = ["url1", "url2", "url3", "url4", "url5"]
+        >>> chunks = chunk_urls(urls, 2)
+        >>> print(chunks)
+        [['url1', 'url2', 'url3'], ['url4', 'url5']]
+    """
+    if num_chunks <= 0:
+        raise ValueError("num_chunks must be positive")
+
+    if not urls:
+        return []
+
+    chunk_size = max(1, len(urls) // num_chunks)
+    chunks = []
+
+    for i in range(0, len(urls), chunk_size):
+        chunk = urls[i : i + chunk_size]
+        if chunk:
+            chunks.append(chunk)
+
+    return chunks
+
+
+__all__ = [
+    "BatchConfig",
+    "ConsolidationConfig",
+    "ShardInfo",
+    "PartitionResult",
+    "BatchResult",
+    "merge_geodataframes",
+    "read_parquet_from_storage",
+    "write_parquet_to_storage",
+    "group_shards_by_partition",
+    "chunk_urls",
+]