earthcatalog 0.2.0__py3-none-any.whl

earthcatalog/workers.py

@@ -0,0 +1,682 @@
+# workers.py
+"""Serializable worker functions for distributed STAC ingestion.
+
+This module provides **module-level functions** that can be pickled and sent
+to remote Dask workers. Unlike the methods in `ingestion_pipeline.py`, these
+functions do not capture `self` and receive all dependencies as explicit parameters.
+
+Architecture:
+    These functions are designed to be called from both:
+    1. ThreadPoolExecutor (local processing)
+    2. Dask distributed workers (cluster processing)
+
+    Each function receives a serialized config dict and storage paths,
+    avoiding closure captures that would prevent pickling.
+
+Key Functions:
+    process_url_batch: Download STAC items from URLs and write shards
+    consolidate_partition: Merge shards for a partition with existing data
+
+Usage:
+    >>> from earthcatalog.workers import process_url_batch
+    >>> from earthcatalog.ingestion_pipeline import ProcessingConfig
+    >>>
+    >>> config = ProcessingConfig(...)
+    >>> config_dict = config.to_dict()
+    >>>
+    >>> # Can be called directly or via Dask
+    >>> result = process_url_batch(
+    ...     urls=["http://example.com/item1.json"],
+    ...     worker_id=0,
+    ...     config_dict=config_dict,
+    ...     job_id="abc-123",
+    ... )
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import tempfile
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import geopandas as gpd
+import pandas as pd
+from tqdm import tqdm
+
+# Conditional async HTTP imports
+try:
+    from .async_http_client import HAS_ASYNC_HTTP, download_stac_items_async
+except ImportError:
+    HAS_ASYNC_HTTP = False
+
+    async def download_stac_items_async(*args, **kwargs) -> list[dict[str, Any]]:  # type: ignore
+        """Dummy async function when async HTTP is not available."""
+        raise ImportError("Async HTTP client not available")
+
+
+from .engines import get_engine
+from .grid_systems import get_grid_system
+from .statistics import IngestionStatistics
+from .storage_backends import get_storage_backend
+
+if TYPE_CHECKING:
+    from .ingestion_pipeline import ProcessingConfig
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Worker Result Types
+# =============================================================================
+
+
+class DownloadResult:
+    """Result of processing a batch of URLs.
+
+    Attributes:
+        shards: List of shard info dictionaries with paths and counts.
+        stats: Statistics from this worker's processing.
+        failed_urls: URLs that failed after all retries.
+    """
+
+    def __init__(
+        self,
+        shards: list[dict[str, Any]],
+        stats: IngestionStatistics,
+        failed_urls: list[str],
+    ):
+        self.shards = shards
+        self.stats = stats
+        self.failed_urls = failed_urls
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize for transmission."""
+        return {
+            "shards": self.shards,
+            "stats": {
+                "urls_processed": self.stats.urls_processed,
+                "urls_failed": self.stats.urls_failed,
+            },
+            "failed_urls": self.failed_urls,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> DownloadResult:
+        """Deserialize from dictionary."""
+        stats = IngestionStatistics()
+        stats_data = data.get("stats", {})
+        # Reconstruct stats from dict
+        if "urls_processed" in stats_data:
+            stats.urls_processed = stats_data["urls_processed"]
+        if "urls_failed" in stats_data:
+            stats.urls_failed = stats_data["urls_failed"]
+        return cls(
+            shards=data.get("shards", []),
+            stats=stats,
+            failed_urls=data.get("failed_urls", []),
+        )
+
+
+class ConsolidationResult:
+    """Result of consolidating a partition.
+
+    Attributes:
+        partition_key: The partition that was consolidated.
+        item_count: Total items in the final partition.
+        existing_count: Items that existed before.
+        new_count: New items added.
+        duplicates_removed: Duplicates that were removed.
+        final_path: Path to the final partition file.
+    """
+
+    def __init__(
+        self,
+        partition_key: str,
+        item_count: int,
+        existing_count: int,
+        new_count: int,
+        duplicates_removed: int = 0,
+        final_path: str = "",
+    ):
+        self.partition_key = partition_key
+        self.item_count = item_count
+        self.existing_count = existing_count
+        self.new_count = new_count
+        self.duplicates_removed = duplicates_removed
+        self.final_path = final_path
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize for transmission."""
+        return {
+            "partition_key": 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,
+        }
+
+
+# =============================================================================
+# Helper Functions
+# =============================================================================
+
+
+def _get_stac_hook(hook_config: str) -> Any:
+    """Get a STAC hook instance from configuration string.
+
+    This is called on workers to instantiate the hook from the serialized config.
+
+    Args:
+        hook_config: Hook configuration string (e.g., "default", "module:pkg:func").
+
+    Returns:
+        Configured hook instance (BaseSTACHook).
+    """
+    from .stac_hooks import get_hook
+
+    return get_hook(hook_config)
+
+
+def _download_stac_item(
+    url: str,
+    timeout: int = 30,
+    retry_attempts: int = 3,
+    hook_config: str = "default",
+) -> dict[str, Any] | None:
+    """Download/generate a single STAC item from URL using configured hook.
+
+    Args:
+        url: URL to the STAC item JSON or data file.
+        timeout: Request timeout in seconds.
+        retry_attempts: Number of retry attempts.
+        hook_config: STAC hook configuration string.
+
+    Returns:
+        Parsed STAC item dictionary, or None if download failed.
+    """
+    hook = _get_stac_hook(hook_config)
+    return hook.fetch(url, timeout=timeout, retry_attempts=retry_attempts)
+
+
+def _download_stac_items_batch(
+    urls: list[str],
+    timeout: int = 30,
+    retry_attempts: int = 3,
+    hook_config: str = "default",
+) -> list[dict[str, Any] | None]:
+    """Download/generate multiple STAC items using configured hook.
+
+    This function leverages the hook's batch processing capability if available,
+    falling back to sequential processing otherwise.
+
+    Args:
+        urls: List of URLs to process.
+        timeout: Request timeout in seconds.
+        retry_attempts: Number of retry attempts.
+        hook_config: STAC hook configuration string.
+
+    Returns:
+        List of STAC items (or None for failures), same order as input.
+    """
+    hook = _get_stac_hook(hook_config)
+    return hook.fetch_batch(urls, timeout=timeout, retry_attempts=retry_attempts)
+
+
+def _get_partition_key(
+    item: dict[str, Any],
+    grid_resolver: Any,
+    temporal_bin: str,
+    enable_global: bool,
+    global_threshold: int,
+) -> str:
+    """Compute the partition key for a STAC item.
+
+    Args:
+        item: STAC item dictionary.
+        grid_resolver: Grid resolver instance.
+        temporal_bin: Temporal binning level ("year", "month", "day").
+        enable_global: Whether to enable global partitioning.
+        global_threshold: Threshold for global partition routing.
+
+    Returns:
+        Partition key string like "h3_82/2024/01" or "global/2024/01".
+    """
+    # Get geometry for grid cell resolution
+    geometry = item.get("geometry")
+
+    # Get grid cells for this item
+    if geometry:
+        tiles, is_spanning = grid_resolver.tiles_for_geometry_with_spanning_detection(geometry)
+        if not tiles:
+            # Fallback to bbox center
+            bbox = item.get("bbox", [0, 0, 0, 0])
+            center_geom = {"type": "Point", "coordinates": [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]}
+            tiles = grid_resolver.tiles_for_geometry(center_geom)
+        grid_cell = tiles[0] if tiles else "unknown"
+
+        # Check for global partition routing
+        if enable_global and len(tiles) > global_threshold:
+            grid_cell = "global"
+    else:
+        # Use bbox center if no geometry
+        bbox = item.get("bbox", [0, 0, 0, 0])
+        center_geom = {"type": "Point", "coordinates": [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]}
+        tiles = grid_resolver.tiles_for_geometry(center_geom)
+        grid_cell = tiles[0] if tiles else "unknown"
+
+    # Get temporal component
+    datetime_str = item.get("properties", {}).get("datetime", "")
+    if datetime_str:
+        try:
+            from datetime import datetime
+
+            dt = datetime.fromisoformat(datetime_str.replace("Z", "+00:00"))
+            if temporal_bin == "year":
+                temporal = str(dt.year)
+            elif temporal_bin == "month":
+                temporal = f"{dt.year}/{dt.month:02d}"
+            else:  # day
+                temporal = f"{dt.year}/{dt.month:02d}/{dt.day:02d}"
+        except (ValueError, AttributeError):
+            temporal = "unknown"
+    else:
+        temporal = "unknown"
+
+    return f"{grid_cell}/{temporal}"
+
+
+def _write_shard_to_scratch(
+    items: list[dict[str, Any]],
+    partition_key: str,
+    scratch_path: str,
+    worker_tag: str,
+    shard_id: int,
+    engine_type: str = "rustac",
+) -> dict[str, Any]:
+    """Write items to a shard file in scratch storage.
+
+    Args:
+        items: STAC items to write.
+        partition_key: Partition key for directory organization.
+        scratch_path: Base scratch directory path.
+        worker_tag: Unique worker identifier.
+        shard_id: Shard number for this worker.
+        engine_type: STAC engine to use.
+
+    Returns:
+        Shard info dictionary with path and metadata.
+    """
+    storage = get_storage_backend(scratch_path)
+
+    # Create shard path
+    safe_partition = partition_key.replace("/", "_")
+    shard_filename = f"{safe_partition}_{worker_tag}_{shard_id}.parquet"
+    shard_path = f"{scratch_path}/shards/{shard_filename}"
+
+    # Ensure directory exists
+    storage.makedirs(Path(shard_path).parent)
+
+    # Convert items to GeoDataFrame using engine
+    engine = get_engine(engine_type)  # type: ignore
+    gdf = engine.items_to_geodataframe(items)
+
+    # Write to scratch
+    with storage.open(shard_path, "wb") as f:
+        gdf.to_parquet(f, index=False, compression="snappy")
+
+    return {
+        "shard_path": shard_path,
+        "partition_key": partition_key,
+        "item_count": len(items),
+        "worker_tag": worker_tag,
+        "shard_id": shard_id,
+    }
+
+
+# =============================================================================
+# Main Worker Functions (Module-Level, Serializable)
+# =============================================================================
+
+
+def process_url_batch(
+    urls: list[str],
+    worker_id: int,
+    config_dict: dict[str, Any],
+    job_id: str,
+    batch_idx: int = 0,
+) -> dict[str, Any]:
+    """Process a batch of URLs, downloading STAC items and writing shards.
+
+    This is a **module-level function** that can be pickled and sent to remote
+    Dask workers. All dependencies are passed as parameters.
+
+    Args:
+        urls: List of STAC item URLs to download.
+        worker_id: Unique worker identifier.
+        config_dict: Serialized ProcessingConfig dictionary.
+        job_id: Job ID for tracking.
+        batch_idx: Batch index for this worker.
+
+    Returns:
+        Dictionary with:
+        - shards: List of shard info dictionaries
+        - stats: Worker statistics as dict
+        - failed_urls: List of URLs that failed
+    """
+    # Reconstruct config from dict
+    from .ingestion_pipeline import ProcessingConfig
+
+    config = ProcessingConfig.from_dict(config_dict)
+
+    # Create worker-local stats
+    worker_stats = IngestionStatistics()
+    failed_urls: list[str] = []
+
+    # Create grid resolver
+    grid_resolver = get_grid_system(
+        config.grid_system,
+        resolution=config.grid_resolution,
+        geojson_path=config.geojson_path if config.grid_system == "geojson" else None,
+    )
+
+    # Unique tag for this worker's shards
+    worker_tag = f"w{worker_id}-{job_id[:8]}-{batch_idx}"
+
+    # Group items by partition as we process
+    partition_items: dict[str, list[dict[str, Any]]] = {}
+    shards: list[dict[str, Any]] = []
+
+    # Choose processing method
+    use_async = config.enable_concurrent_http and HAS_ASYNC_HTTP and len(urls) >= config.batch_size
+
+    if use_async:
+        # Async HTTP processing
+        items = _download_batch_async(urls, config)
+        for item in items:
+            if item:
+                worker_stats.record_url_processed(success=True)
+                partition_key = _get_partition_key(
+                    item,
+                    grid_resolver,
+                    config.temporal_bin,
+                    config.enable_global_partitioning,
+                    config.global_partition_threshold,
+                )
+                if partition_key not in partition_items:
+                    partition_items[partition_key] = []
+                partition_items[partition_key].append(item)
+    else:
+        # Synchronous processing with progress bar
+        # Get hook config from ProcessingConfig
+        hook_config = config_dict.get("stac_hook", "default")
+        with tqdm(total=len(urls), desc=f"Worker {worker_id}", unit="urls", leave=False) as pbar:
+            for url in urls:
+                fetched_item = _download_stac_item(
+                    url,
+                    timeout=config.request_timeout,
+                    retry_attempts=config.retry_attempts,
+                    hook_config=hook_config,
+                )
+                if fetched_item:
+                    worker_stats.record_url_processed(success=True)
+                    partition_key = _get_partition_key(
+                        fetched_item,
+                        grid_resolver,
+                        config.temporal_bin,
+                        config.enable_global_partitioning,
+                        config.global_partition_threshold,
+                    )
+                    if partition_key not in partition_items:
+                        partition_items[partition_key] = []
+                    partition_items[partition_key].append(fetched_item)
+                else:
+                    worker_stats.record_url_processed(success=False)
+                    failed_urls.append(url)
+                pbar.update(1)
+
+    # Write shards for each partition
+    shard_counter = 0
+    for partition_key, items in partition_items.items():
+        # Split into shard-sized chunks
+        for i in range(0, len(items), config.items_per_shard):
+            chunk = items[i : i + config.items_per_shard]
+            shard_info = _write_shard_to_scratch(
+                items=chunk,
+                partition_key=partition_key,
+                scratch_path=config.scratch_location,
+                worker_tag=worker_tag,
+                shard_id=shard_counter,
+                engine_type=config.stac_engine,
+            )
+            shards.append(shard_info)
+            shard_counter += 1
+
+    logger.info(
+        f"Worker {worker_id}: processed {len(urls)} URLs, wrote {len(shards)} shards, {len(failed_urls)} failures"
+    )
+
+    # Serialize stats as simple dict with basic counters
+    stats_dict = {
+        "urls_processed": worker_stats.urls_processed,
+        "urls_failed": worker_stats.urls_failed,
+    }
+
+    return {
+        "shards": shards,
+        "stats": stats_dict,
+        "failed_urls": failed_urls,
+    }
+
+
+def _download_batch_async(urls: list[str], config: ProcessingConfig) -> list[dict[str, Any]]:
+    """Download URLs using async HTTP."""
+    try:
+        # Try to get or create event loop
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # Running in async context (Jupyter, etc.)
+                with ThreadPoolExecutor(max_workers=1) as _:
+                    try:
+                        import nest_asyncio
+
+                        nest_asyncio.apply()
+                    except ImportError:
+                        pass  # nest_asyncio not available
+                    return loop.run_until_complete(
+                        download_stac_items_async(
+                            urls,
+                            concurrent_requests=config.concurrent_requests,
+                            connection_pool_size=config.connection_pool_size,
+                            request_timeout=config.request_timeout,
+                            retry_attempts=config.retry_attempts,
+                            retry_delay=config.retry_delay,
+                            batch_size=config.batch_size,
+                        )
+                    )
+        except RuntimeError:
+            pass
+
+        # Create new event loop
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            return loop.run_until_complete(
+                download_stac_items_async(
+                    urls,
+                    concurrent_requests=config.concurrent_requests,
+                    connection_pool_size=config.connection_pool_size,
+                    request_timeout=config.request_timeout,
+                    retry_attempts=config.retry_attempts,
+                    retry_delay=config.retry_delay,
+                    batch_size=config.batch_size,
+                )
+            )
+        finally:
+            loop.close()
+    except (ValueError, TypeError, RuntimeError, ConnectionError) as e:
+        logger.error(f"Async download failed: {e}, falling back to sync")
+        # Fallback to sync
+        items = []
+        for url in urls:
+            item = _download_stac_item(url, config.request_timeout, config.retry_attempts)
+            if item:
+                items.append(item)
+        return items
+
+
+def consolidate_partition(
+    partition_key: str,
+    shard_paths: list[str],
+    config_dict: dict[str, Any],
+) -> dict[str, Any]:
+    """Consolidate shards for a partition, merging with existing data.
+
+    This is a **module-level function** that can be pickled and sent to remote
+    Dask workers. All dependencies are passed as parameters.
+
+    Deduplication uses newer datetime wins strategy: when items have the same ID,
+    the one with the more recent datetime property is kept.
+
+    Args:
+        partition_key: Partition key (e.g., "h3_82/2024/01").
+        shard_paths: List of shard file paths to merge.
+        config_dict: Serialized ProcessingConfig dictionary.
+
+    Returns:
+        Dictionary with consolidation results:
+        - partition_key: The partition that was consolidated
+        - item_count: Total items in final partition
+        - existing_count: Items that existed before
+        - new_count: New items added
+        - duplicates_removed: Number of duplicates removed
+        - final_path: Path to final partition file
+    """
+    from .ingestion_pipeline import ProcessingConfig
+
+    config = ProcessingConfig.from_dict(config_dict)
+
+    # Get storage backends
+    catalog_storage = get_storage_backend(config.output_catalog)
+    scratch_storage = get_storage_backend(config.scratch_location)
+
+    # Compute final path
+    final_path = f"{config.output_catalog}/{partition_key}/data.parquet"
+
+    existing_count = 0
+    all_gdfs: list[gpd.GeoDataFrame] = []
+
+    # Use temp directory for staging
+    with tempfile.TemporaryDirectory(dir=config.temp_dir_location) as temp_dir:
+        temp_existing_path = Path(temp_dir) / "existing.parquet"
+        temp_merged_path = Path(temp_dir) / "merged.parquet"
+
+        # Step 1: Download existing partition if it exists
+        if catalog_storage.exists(final_path):
+            try:
+                logger.debug(f"Partition {partition_key}: downloading existing data")
+                with catalog_storage.open(final_path, "rb") as f:
+                    binary_data = f.read()
+                    # Write to temp file so we can use gpd.read_parquet
+                    with open(temp_existing_path, "wb") as temp_f:
+                        temp_f.write(binary_data)
+                    existing_gdf = gpd.read_parquet(temp_existing_path)
+                    existing_count = len(existing_gdf)
+                    all_gdfs.append(existing_gdf)
+            except (OSError, ValueError, TypeError) as e:
+                logger.error(f"Error reading existing partition {final_path}: {e}")
+
+        # Step 2: Read all new shards
+        for idx, shard_path in enumerate(shard_paths):
+            try:
+                with scratch_storage.open(shard_path, "rb") as f:
+                    binary_data = f.read()
+                    # Write to temp file so we can use gpd.read_parquet
+                    temp_shard_path = Path(temp_dir) / f"shard_{idx}.parquet"
+                    with open(temp_shard_path, "wb") as temp_f:
+                        temp_f.write(binary_data)
+                    gdf = gpd.read_parquet(temp_shard_path)
+                    all_gdfs.append(gdf)
+            except (OSError, ValueError, TypeError) as e:
+                logger.error(f"Error reading shard {shard_path}: {e}")
+                continue
+
+        if not all_gdfs:
+            return {
+                "partition_key": partition_key,
+                "item_count": 0,
+                "existing_count": 0,
+                "new_count": 0,
+                "duplicates_removed": 0,
+                "final_path": final_path,
+            }
+
+        # Step 3: Merge all data
+        merged = pd.concat(all_gdfs, ignore_index=True)
+        original_count = len(merged)
+
+        # Step 4: Deduplicate - newer datetime wins
+        if "datetime" in merged.columns and "id" in merged.columns:
+            # Sort by datetime descending so newer items come first
+            try:
+                merged["_dt_sort"] = pd.to_datetime(merged["datetime"], errors="coerce")
+                merged = merged.sort_values("_dt_sort", ascending=False, na_position="last")
+                merged = merged.drop_duplicates(subset=["id"], keep="first")
+                merged = merged.drop(columns=["_dt_sort"])
+            except (ValueError, TypeError):
+                # Fallback to simple dedup
+                merged = merged.drop_duplicates(subset=["id"], keep="last")
+        elif "id" in merged.columns:
+            merged = merged.drop_duplicates(subset=["id"], keep="last")
+
+        duplicates_removed = original_count - len(merged)
+
+        # Step 5: Sort by configured key
+        if config.sort_key in merged.columns:
+            merged = merged.sort_values(config.sort_key, ascending=config.sort_ascending)
+
+        # Step 6: Write to temp file
+        # Ensure geometry column is properly set
+        if "geometry" in merged.columns:
+            merged_gdf = gpd.GeoDataFrame(merged, geometry="geometry")
+        else:
+            merged_gdf = gpd.GeoDataFrame(merged)
+        merged_gdf.to_parquet(temp_merged_path, index=False, compression="snappy")
+
+        # Step 7: Upload to final location
+        catalog_storage.makedirs(Path(final_path).parent)
+        if final_path.startswith("s3://"):
+            catalog_storage.upload(str(temp_merged_path), final_path)
+        else:
+            # Local storage - atomic move
+            import shutil
+
+            shutil.copy2(temp_merged_path, final_path)
+
+        new_count = len(merged) - existing_count
+
+    logger.info(
+        f"Partition {partition_key}: {existing_count} existing + {new_count} new = "
+        f"{len(merged)} total ({duplicates_removed} duplicates removed)"
+    )
+
+    return {
+        "partition_key": partition_key,
+        "item_count": len(merged),
+        "existing_count": existing_count,
+        "new_count": new_count,
+        "duplicates_removed": duplicates_removed,
+        "final_path": final_path,
+    }
+
+
+__all__ = [
+    "DownloadResult",
+    "ConsolidationResult",
+    "process_url_batch",
+    "consolidate_partition",
+]