Pipeline API¶

def search(
    self,
    provider: SearchProvider | None = None,
    aoi: BaseGeometry | dict[str, Any] | str | Path | None = None,
    **search_kwargs: Any,
) -> GeoDataFrame[AssetSchema]:
    """Execute a search.

    Uses *provider* when given; otherwise falls back to the job's configured
    ``search_provider``. Runtime search parameters win over the job's fixed
    ``target_aoi``. The resolved AOI is passed to the provider as
    ``intersects``.

    Args:
        provider: Optional search provider to execute. Defaults to the
            provider configured on the job.
        aoi: Optional AOI geometry overriding ``job.target_aoi``.
        **search_kwargs: Additional arguments used to update the provider
            before execution (e.g. ``start_datetime``, ``end_datetime``).

    Returns:
        A validated GeoDataFrame of matched assets.

    Raises:
        ValueError: If no provider is given and none is configured on the job.
    """
    provider = provider or self.search_provider
    if provider is None:
        raise ValueError(
            "No search provider configured. Pass one or set it in the job config."
        )

    logger.info(
        "search_called",
        provider=_callable_name(provider),
    )

    resolved_aoi = (
        normalize_geometry_input(aoi)
        if aoi is not None
        else self.effective_target_aoi
    )
    if resolved_aoi is not None:
        search_kwargs["intersects"] = resolved_aoi

    bound_provider = update_callable(provider, **search_kwargs)

    return bound_provider()

def build_tasks(
    self,
    assets: GeoDataFrame[AssetSchema],
    task_builder: TaskBuilder | None = None,
    **builder_kwargs: Any,
) -> Sequence[ExtractionTask]:
    """Build extraction tasks from search results.

    Args:
        assets: GeoDataFrame of assets returned by a search provider.
        task_builder: Optional task builder used to group assets into tasks.
            Defaults to the builder configured on the job.
        **builder_kwargs: Additional arguments used to update the task
            builder before execution (e.g. ``cells_per_task``).

    Returns:
        A sequence of prepared ``ExtractionTask`` objects.

    Raises:
        ValueError: If no builder is given and none is configured on the job.
    """
    task_builder = task_builder or self.task_builder
    if task_builder is None:
        raise ValueError(
            "No task builder configured. Pass one or set it in the job config."
        )

    if assets.empty:
        return []

    logger.info(
        "build_tasks_start",
        builder=_callable_name(task_builder),
        assets=len(assets),
    )

    if builder_kwargs:
        task_builder = update_callable(task_builder, **builder_kwargs)

    assert task_builder is not None
    return task_builder(assets, self)

def execute(
    self,
    tasks: Sequence[ExtractionTask],
    executor: Executor | None = None,
) -> GeoDataFrame[ArtifactSchema]:
    """Run prepared tasks and return the combined artifacts.

    Args:
        tasks: Extraction tasks to execute.
        executor: Optional executor. Defaults to ``LocalExecutor()``.

    Returns:
        A validated ``GeoDataFrame[ArtifactSchema]``.
    """
    if not tasks:
        return cast(
            GeoDataFrame[ArtifactSchema], ArtifactSchema.empty_geodataframe()
        )

    selected_executor = executor or LocalExecutor()
    logger.info(
        "execute_start",
        task_count=len(tasks),
        executor=selected_executor.__class__.__name__,
    )
    return selected_executor(tasks)

def write_catalog(
    self,
    artifacts: GeoDataFrame[ArtifactSchema],
    uri: str | Path | None = None,
) -> str:
    """Write the artifact catalog to parquet.

    Args:
        artifacts: GeoDataFrame of extracted artifacts.
        uri: Destination URI. Defaults to
            ``{output_uri}/artifacts.parquet``.

    Returns:
        The URI where the catalog was written.
    """
    if uri is None:
        uri = f"{self.output_uri.rstrip('/')}/artifacts.parquet"

    str_uri = str(uri)
    if not str_uri.startswith("s3://"):
        path = Path(str_uri.removeprefix("file://"))
        path.parent.mkdir(parents=True, exist_ok=True)

    artifacts.to_parquet(str_uri)  # pyright: ignore[reportArgumentType]
    return str_uri

@classmethod
def load_from_config(
    cls,
    config_dir: str | Path,
    config_name: str = "main_config",
    overrides: list[str] | None = None,
) -> ExtractionJob:
    """Load and validate an ``ExtractionJob`` from a Hydra config package.

    This is the recommended way to consume a Hydra config package: it
    initializes the config directory, composes the configuration,
    recursively instantiates all plugins, and validates the result.

    Args:
        config_dir: Directory containing the Hydra config package.
        config_name: Name of the root config file (without ``.yaml``).
        overrides: Optional Hydra command-line style overrides, e.g.
            ``["grid_dist=grid_50km"]``.

    Returns:
        A validated ``ExtractionJob`` instance.

    Example::

        from aereo.pipeline import ExtractionJob

        job = ExtractionJob.load_from_config(
            "examples/config_package",
            overrides=["grid_dist=grid_50km"],
        )
    """
    from hydra import compose, initialize_config_dir
    from omegaconf import OmegaConf

    config_dir = Path(config_dir).resolve()
    with initialize_config_dir(version_base=None, config_dir=str(config_dir)):
        cfg = compose(config_name=config_name, overrides=overrides or [])
        plain_cfg = OmegaConf.to_container(cfg, resolve=True)
        plain_cfg = _strip_unknown_job_keys(plain_cfg, cls)
        prepared_cfg = _prepare_config_for_instantiate(plain_cfg)
        instantiated = hydra.utils.instantiate(prepared_cfg, _convert_="all")

    return cls._from_instantiated(
        instantiated, f"config at {config_dir}/{config_name}"
    )

@classmethod
def from_yaml(cls, path: str | Path) -> ExtractionJob:
    """Load an ExtractionJob from a single YAML file using Hydra.

    Loads the configuration via OmegaConf, recursively instantiates all
    target classes using hydra.utils.instantiate, and returns an
    ExtractionJob instance.

    The expected YAML layout places ``grid_dist``, ``output_uri``,
    ``read`` and ``write`` as top-level keys, enabling Hydra config package
    composition::

        defaults:
          - grid_dist: default
          - read: sentinel2
          - write: geotiff
          - _self_

        output_uri: /tmp/extraction

    Args:
        path: Path to the YAML config file.

    Returns:
        A validated ``ExtractionJob`` instance.

    Raises:
        ValueError: If Hydra instantiation does not produce an
            ``ExtractionJob`` or a dict.
    """
    from omegaconf import OmegaConf

    path = Path(path)
    cfg = OmegaConf.load(path)

    plain_cfg = OmegaConf.to_container(cfg, resolve=True)
    plain_cfg = _strip_unknown_job_keys(plain_cfg, cls)
    prepared_cfg = _prepare_config_for_instantiate(plain_cfg)
    instantiated = hydra.utils.instantiate(prepared_cfg, _convert_="all")
    return cls._from_instantiated(instantiated, f"configuration at {path}")