nirs4all.pipeline.storage.array_store module

Parquet-based storage for prediction arrays.

Stores dense numerical arrays (y_true, y_pred, y_proba, optional replay X, sample_indices, weights) in one Parquet file per dataset, with Zstd compression. Each file embeds lightweight metadata columns (model_name, fold_id, partition, metric, val_score, task_type) making it self-describing and portable.

The relational metadata (runs, pipelines, chains, full scores) stays in SQLite. This module only manages the array sidecar files.

Classification string/int64 targets use optional versioned JSON label columns, with a sorted typed vocabulary and integer positions. These positions describe storage only: they never change the model’s probability-column class order. Purely numerical files retain their historical schema and target representation.

Workspace layout:

workspace/
    store.sqlite          # Metadata only
    arrays/
        wheat.parquet     # All prediction arrays for dataset "wheat"
        corn.parquet      # All prediction arrays for dataset "corn"
        _tombstones.json  # Pending deletes (prediction_ids)

class nirs4all.pipeline.storage.array_store.ArrayStore(base_dir: Path, *, read_only: bool = False)

Bases: object

Parquet-backed storage for prediction arrays.

Arrays live under base_dir / 'arrays', one .parquet file per dataset. Writes append row groups; deletes use a tombstone file that is applied during compact().

Parameters:

base_dir – Workspace root directory. The arrays/ subdirectory is created automatically.

property arrays_dir: Path

Directory containing the Parquet array files.

compact(dataset_name: str | None = None, live_ids: set[str] | None = None) → dict[str, dict[str, Any]]

Rewrite Parquet file(s): apply tombstones, deduplicate, re-sort.

Parameters:

• dataset_name – If given, compact only that dataset’s file. If None, compact all datasets.

• live_ids – Prediction IDs that are LIVE in committed SQLite metadata. When given, a tombstoned id that is still live is NEVER removed — its stale tombstone is dropped instead. Stale tombstones arise from a crash or rollback between a SQLite write and the matching tombstone write (e.g. the save_prediction upsert window). When None, tombstones are applied as-is; callers with access to the SQLite metadata should prefer WorkspaceStore.compact_arrays, which also snapshots the live ids under the process lock.

Returns:

Stats per dataset:

{dataset: {rows_before, rows_after, rows_removed,
           bytes_before, bytes_after}}

delete_batch(prediction_ids: set[str], dataset_name: str | None = None) → int

Mark prediction_ids as deleted (tombstone).

Physical removal happens during compact().

Parameters:

• prediction_ids – Set of prediction IDs to mark as deleted.

• dataset_name – Optional dataset filter.

Returns:

Number of IDs marked for deletion.

delete_dataset(dataset_name: str) → bool

Delete the entire Parquet file for a dataset.

Parameters:

dataset_name – Name of the dataset to delete.

Returns:

True if the file existed and was deleted.

has_pending_tombstones() → bool

Return True when deletes are pending physical removal.

integrity_check(expected_ids: set[str] | None = None) → dict[str, Any]

Check Parquet health.

Parameters:

expected_ids – If given, cross-check against these prediction IDs.

Returns:

{orphan_ids: [...], missing_ids: [...], corrupt_files: [...]}

list_datasets() → list[str]

Return dataset names that have Parquet files.

load_batch(prediction_ids: list[str], dataset_name: str | None = None) → dict[str, dict[str, Any]]

Load arrays for multiple predictions.

Uses predicate pushdown on prediction_id for efficient reads.

Parameters:

• prediction_ids – List of prediction IDs to load.

• dataset_name – If given, reads only that dataset’s file.

Returns:

{y_true, y_pred, y_proba, X, spectra, sample_indices, weights, sample_metadata, result_metadata}}``

Return type:

``{prediction_id

load_dataset(dataset_name: str) → DataFrame

Load the full Parquet file for a dataset.

Returns all columns (arrays + metadata). This is the ‘portable read’ — the returned DataFrame is self-describing.

Parameters:

dataset_name – Name of the dataset.

Returns:

A Polars DataFrame with all rows and columns.

Raises:

FileNotFoundError – If no Parquet file exists for this dataset.

load_single(prediction_id: str, dataset_name: str | None = None) → dict[str, Any] | None

Load arrays for one prediction.

Parameters:

• prediction_id – Prediction ID to load.

• dataset_name – If given, reads only that dataset’s file.

Returns:

Dict with array fields, or None if not found.

pending_tombstone_count() → int

Number of deletes pending physical removal.

Cheap (single small JSON read, no Parquet access) — used to gate the startup reconciliation and the auto-compaction in WorkspaceStore.

save_batch(records: list[dict]) → int

Write prediction arrays to Parquet.

Each dict should contain:

{prediction_id, dataset_name, model_name, fold_id, partition,
 metric, val_score, task_type, y_true, y_pred, y_proba,
 y_proba_shape, X, X_shape, spectra, spectra_shape,
 sample_indices, weights, sample_metadata, result_metadata}

Records are grouped by dataset_name; each group is appended to its dataset Parquet file. Writes are idempotent by prediction_id (duplicates resolved on compact()).

Parameters:

records – List of prediction array records.

Returns:

Number of rows written.

stats() → dict[str, Any]

Return storage stats.

Returns:

``{total_files, total_rows, total_bytes,

datasets: {name: {rows, bytes}}}``