Configuration¶

get_dataset ¶

get_dataset(dataset_id: DatasetId) -> Dataset

Get a registered dataset bundle by ID.

Parameters:

dataset_id (DatasetId) –

Dataset registry id, such as "pozzato-2022".

Returns:

Dataset –

Registered dataset bundle with ingest and normalize helpers.

Raises:

ValueError –

If the dataset id is unknown.

Examples:

>>> dataset = get_dataset("pozzato-2022")
>>> dataset.spec.dataset_id
'pozzato-2022'

DatasetInfo `dataclass` ¶

DatasetInfo(
    name: str | None = None,
    year: int | None = None,
    author: str | None = None,
    parent_dataset_id: DatasetId | None = None,
    misc: dict[str, object] = dict(),
    description: str | None = None,
)

Human-readable dataset metadata used by registry entries.

Attributes:

name (str | None) –

Display name.
year (int | None) –

Publication or release year.
author (str | None) –

Primary author or source organization.
parent_dataset_id (DatasetId | None) –

Registry id of the source dataset when this dataset is derived from another one.
misc (dict[str, object]) –

Additional source-specific metadata.
description (str | None) –

Optional longer description.

DatasetSpec `dataclass` ¶

DatasetSpec(
    dataset_id: DatasetId,
    dataset_type: DatasetTypeId,
    info: DatasetInfo | None = None,
    processing_stages: dict[
        DatasetStageId, STAGE_SPECS
    ] = dict(),
)

Storage layout and processing-stage configuration for one dataset.

processing_stages maps stage ids to ingest and normalize specs. Store paths are relative; the configured DataProcessingStore resolves them to local or remote locations.

Attributes:

dataset_id (DatasetId) –

Stable registry and storage id.
dataset_type (DatasetTypeId) –

Dataset origin/type label.
info (DatasetInfo | None) –

Optional human-readable metadata.
processing_stages (dict[DatasetStageId, STAGE_SPECS]) –

Stage configuration keyed by stage id.

Examples:

>>> from batgrad.contracts.mapping import DatasetStageId, DatasetTypeId
>>> spec = DatasetSpec(
...     dataset_id="my-dataset",
...     dataset_type=DatasetTypeId.published,
... )
>>> spec.manifest(DatasetStageId.ingested)
'type=published/dataset=my-dataset/source=ingested/manifest.parquet'

Methods:

source_root –

Relative root for one dataset stage.
source_file –

Relative file path inside one dataset stage.
manifest –

Relative manifest path for one dataset stage.

root `property` ¶

root: str

Relative root location of the dataset.

Absolute root location will be resolved by DataStore.

Returns:

str –

Relative dataset root path.

source_root ¶

source_root(source: DatasetStageId) -> str

Relative root for one dataset stage.

Parameters:

source (DatasetStageId) –

Stage id.

Returns:

str –

Relative store root for files generated by that stage.

source_file ¶

source_file(source: DatasetStageId, file_name: str) -> str

Relative file path inside one dataset stage.

Parameters:

source (DatasetStageId) –

Stage id.
file_name (str) –

File name relative to the stage root.

Returns:

str –

Relative store path for the stage file.

manifest ¶

manifest(source: DatasetStageId) -> str

Relative manifest path for one dataset stage.

Parameters:

source (DatasetStageId) –

Stage id.

Returns:

str –

Relative path to the stage manifest parquet file.

Dataset `dataclass` ¶

Dataset(
    spec: DatasetSpec,
    raw_adapter: RawDatasetAdapter | None = None,
)

Registered dataset bundle and main processing interface.

Dataset combines a DatasetSpec with an optional raw adapter. Use it when running standard ingest/normalize workflows from registry entries. A dataset without raw_adapter cannot ingest raw files but may still support normalize or interactive loading from existing manifests.

Attributes:

spec (DatasetSpec) –

Dataset storage layout and stage configuration.
raw_adapter (RawDatasetAdapter | None) –

Optional raw-file adapter used by ingest.

Examples:

>>> dataset = get_dataset("pozzato-2022")
>>> dataset.spec.dataset_id
'pozzato-2022'

Methods:

ingest –

Run this dataset's raw-to-ingested parquet stage.
normalize –

Run this dataset's ingested-to-normalized parquet stage.
normalize_interactive –

Run selected normalization work into scratch storage for exploration.
load_interactive –

Load selected rows from an existing stage manifest as an interactive run.
load_interactive_manifest –

Wrap an already-loaded manifest as an interactive run.

ingest ¶

ingest(
    input_store: DataProcessingStore,
    output_store: DataProcessingStore,
    config: IngestStageConfig,
    *,
    scratch_store: DataProcessingStore | None = None,
    tasks: tuple[IngestTask, ...] | None = None,
) -> None

Run this dataset's raw-to-ingested parquet stage.

Parameters:

input_store (DataProcessingStore) –

Store containing raw source files.
output_store (DataProcessingStore) –

Store receiving ingested shards and manifest.
config (IngestStageConfig) –

Ingest runtime and parquet-writing settings.
scratch_store (DataProcessingStore | None, default: None ) –

Optional temporary-output store. Defaults to output_store; use a separate scratch store when final output should not hold intermediate task files.
tasks (tuple[IngestTask, ...] | None, default: None ) –

Optional task subset for retries, debugging, or tests. When omitted, the raw adapter plans tasks from source files.

Raises:

TypeError –

If this dataset has no raw adapter.

normalize ¶

normalize(
    input_store: DataProcessingStore,
    output_store: DataProcessingStore,
    config: NormalizeStageConfig,
    *,
    scratch_store: DataProcessingStore | None = None,
    tasks: tuple[NormalizeTask, ...] | None = None,
    dry_run: bool = False,
) -> None

Run this dataset's ingested-to-normalized parquet stage.

Parameters:

input_store (DataProcessingStore) –

Store containing ingested manifest and shards.
output_store (DataProcessingStore) –

Store receiving normalized shards and manifest.
config (NormalizeStageConfig) –

Normalize runtime, validation, and parquet-writing settings.
scratch_store (DataProcessingStore | None, default: None ) –

Optional temporary-output store. Defaults to output_store.
tasks (tuple[NormalizeTask, ...] | None, default: None ) –

Optional pre-planned task subset for retries, debugging, or tests. When omitted, tasks are planned from the ingested manifest.
dry_run (bool, default: False ) –

Validate task preparation and checks without writing final normalized outputs.

normalize_interactive ¶

normalize_interactive(
    input_store: DataProcessingStore,
    scratch_store: DataProcessingStore,
    config: NormalizeStageConfig,
    *,
    protocols: object = None,
    group_values: object = None,
    annotate: bool = True,
    source_run: InteractiveStageRun | None = None,
    normalize_spec: NormalizeStageSpec | None = None,
) -> InteractiveStageRun

Run selected normalization work into scratch storage for exploration.

Parameters:

input_store (DataProcessingStore) –

Store containing ingested or source interactive data.
scratch_store (DataProcessingStore) –

Store receiving temporary interactive outputs.
config (NormalizeStageConfig) –

Normalize runtime, validation, and parquet-writing settings.
protocols (object, default: None ) –

Optional protocol selector, such as one protocol id or a list of protocol ids.
group_values (object, default: None ) –

Optional task group selector or list of selectors. Keys are group column names such as cell id or cycle index.
annotate (bool, default: True ) –

Whether public annotation columns are included in outputs.
source_run (InteractiveStageRun | None, default: None ) –

Optional prior unresampled normalized interactive run to reuse when trying resampling changes.
normalize_spec (NormalizeStageSpec | None, default: None ) –

Optional normalization spec override.

Returns:

InteractiveStageRun –

Interactive run handle for scanning outputs and cleaning scratch
InteractiveStageRun –

files.

load_interactive ¶

load_interactive(
    input_store: DataProcessingStore,
    *,
    source: DatasetStageId | str,
    protocols: object = None,
    group_values: object = None,
) -> InteractiveStageRun

Load selected rows from an existing stage manifest as an interactive run.

Parameters:

input_store (DataProcessingStore) –

Store containing the stage manifest and data shards.
source (DatasetStageId | str) –

Stage id to load, such as DatasetStageId.ingested.
protocols (object, default: None ) –

Optional protocol selector.
group_values (object, default: None ) –

Optional task group selector or list of selectors.

Returns:

InteractiveStageRun –

Interactive run handle over selected existing stage data.

load_interactive_manifest ¶

load_interactive_manifest(
    input_store: DataProcessingStore,
    *,
    source: DatasetStageId | str,
    manifest: DataFrame,
    protocols: object = None,
    group_values: object = None,
) -> InteractiveStageRun

Wrap an already-loaded manifest as an interactive run.

Parameters:

input_store (DataProcessingStore) –

Store containing the stage data shards referenced by the manifest.
source (DatasetStageId | str) –

Stage id the manifest belongs to.
manifest (DataFrame) –

Manifest rows to expose interactively.
protocols (object, default: None ) –

Optional protocol selector used for protocol order.
group_values (object, default: None ) –

Optional task group selector retained on the run.

Returns:

InteractiveStageRun –

Interactive run handle over the provided manifest rows.

InteractiveStageRun `dataclass` ¶

InteractiveStageRun(
    output_store: DataProcessingStore,
    stage_spec: InteractiveStageSpec,
    segment_col: MappingSpec,
    protocol_order: tuple[str, ...] = (),
    dataset_spec: DatasetSpec | None = None,
    input_store: DataProcessingStore | None = None,
    protocols: object = None,
    group_values: object = None,
    source_stage: DatasetStageId | None = None,
    manifest_path: str | None = None,
    manifest_frame: DataFrame | None = None,
    run_root: str | None = None,
)

Handle for selected existing or scratch-generated stage outputs.

Interactive runs are returned by Dataset.normalize_interactive and Dataset.load_interactive. They expose a selected manifest, lazy scanning of referenced data segments, and cleanup for scratch-backed runs.

Attributes:

output_store (DataProcessingStore) –

Store containing the selected output segments.
stage_spec (InteractiveStageSpec) –

Stage configuration used to interpret protocols and outputs.
segment_col (MappingSpec) –

Manifest column containing parquet segment references.
protocol_order (tuple[str, ...]) –

Preferred protocol ordering for UI or plotting callers.
dataset_spec (DatasetSpec | None) –

Dataset configuration, when available.
input_store (DataProcessingStore | None) –

Source store used to create the run, when available.
protocols (object) –

Protocol selector used to create the run.
group_values (object) –

Group selector used to create the run.
source_stage (DatasetStageId | None) –

Stage represented by the run.
manifest_path (str | None) –

Manifest path for scratch-backed runs.
manifest_frame (DataFrame | None) –

Already-loaded manifest for selected existing data.
run_root (str | None) –

Scratch run root cleaned by clean.

Methods:

manifest –

Load or return the manifest rows represented by this run.
protocol_spec –

Return the stage protocol spec for a protocol id or value.
scan –

Scan all data segments referenced by the run manifest.
clean –

Delete this run's scratch root when it owns one.
iter_sources –

Iterate manifest rows with segment readers for each row.

manifest ¶

manifest() -> DataFrame

Load or return the manifest rows represented by this run.

Returns:

DataFrame –

Manifest dataframe for the selected rows.

protocol_spec ¶

protocol_spec(protocol: object) -> InteractiveProtocolSpec

Return the stage protocol spec for a protocol id or value.

Parameters:

protocol (object) –

Protocol id or enum-like value.

Returns:

InteractiveProtocolSpec –

Matching protocol spec from the stage configuration.

scan ¶

scan() -> LazyFrame

Scan all data segments referenced by the run manifest.

Returns:

LazyFrame –

Lazy frame over the selected output segments.

Raises:

FileNotFoundError –

If the selected manifest contains no segments.

clean ¶

clean() -> None

Delete this run's scratch root when it owns one.

iter_sources ¶

iter_sources() -> Iterator[
    tuple[dict[str, object], SegmentSource]
]

Iterate manifest rows with segment readers for each row.

Returns:

Iterator[tuple[dict[str, object], SegmentSource]] –

Iterator of manifest row dictionaries and segment sources.

Configuration¶

get_dataset ¶

DatasetInfo dataclass ¶

DatasetSpec dataclass ¶

root property ¶

source_root ¶

source_file ¶

manifest ¶

Dataset dataclass ¶

ingest ¶

normalize ¶

normalize_interactive ¶

load_interactive ¶

load_interactive_manifest ¶

InteractiveStageRun dataclass ¶

manifest ¶

protocol_spec ¶

scan ¶

clean ¶

iter_sources ¶

DatasetInfo `dataclass` ¶

DatasetSpec `dataclass` ¶

root `property` ¶

Dataset `dataclass` ¶

InteractiveStageRun `dataclass` ¶