StorageSession

Create

session = StorageSession.from_plan(
    plan,  # StoragePlan from compile / parse_storage_plan
    run_id=...,
    params={"user_id": ..., "project_id": ...},
)

Tenant paths: user_id / project_id → object keys under the binding root.

Write path

session.write(dataset, rows, shard_id=…)
  → output_spec.backend
  → write_backend_for(backend).write(…)
  → parquet buffer / PG upsert / Qdrant upsert

write(..., run_id=…) currently accepts run_id for signature compatibility and ignores it; rows always land under session.run_id.

Flush / drain on close iterates all registered write backends (builtins are lazy-registered).

Early shuffle (parquet outputs)

When the dataset declares DRTML shuffle: { by, buckets }, part rotation splits each flush table by stable blake2 buckets (storage/dataset/shuffle.py):

• object keys under bucket-NNN/part-….parquet;
• manifest shard ids like b017-w0-0.

Shuffle parts grow to max_mb_per_part (or the flush MB/row limit). The bucket ParquetWriter stays open across flushes until that limit — closing after every flush made the part-size cap a no-op. Rotation still happens at the size/row limit; remaining writers close on session finalize. Prefer moderate shuffle.buckets: peak open writers tracks touched buckets that are still filling a part.

There is no post-hoc compact after the run: driver merge only writes the shard manifest. Parts stay as they were assembled while writing (part-w* from workers).

Optional param storage.shuffle_buckets (applied at session open) overrides shuffle.buckets on every output that already has shuffle:.

See Datasets and reduce_merge.

Read path

• read_batches — parquet-only (RecordBatch).
• read_run_rows / iter_run_rows — via read backend (parquet/pg/qdrant).
• Parquet logical-row helper: storage.run.session_read.iter_logical_rows → row_materialize (not a StorageSession method).

Streaming parquet rows

API	Use
session_read.iter_parquet_run_rows(..., filters=, group_by=, sum_columns=)	Hot path — stream filtered logical rows; optional incremental SUM over group_by
session_read.read_parquet_run_rows(...)	Materializes list(iter_…); small fixtures / tests only

Without group_by, peak RAM tracks one read batch (parquet can slice by row group). With group_by + sum_columns (both required together), peak tracks unique filtered keys. Do not load a full run into a Python list on production paths.

iter_run_rows dispatches through backend_registry; the parquet adapter uses the streaming iterator.

Vector batches (geometry windows)

For cluster fit/assign (and similar geometry), storage/run/vector_batch_iter.py streams bounded numpy windows: vectors, row ids, optional metas. It sits on top of logical-row reads — not a new dataset backend.

Callers pass column names, equality filters, optional shard id, and optional max_rows (probe sample cap). Peak tracks batch_rows, not the full run. Do not confuse this with feed ProcessorBatch, Arrow RecordBatch, or embed train pair batches.

See Streaming cluster.

Checkpoint

Cadence from datasets.*.checkpoint (column from feed input, not vocab) → partition_loop glue.
flush.boundary_columns — parquet only; postgres/qdrant → validate error.

Part rotation

run/part_rotation.py — part-*.parquet by MB/rows/boundary, optionally per shuffle bucket.
Shuffle writers stay open until the part-size limit (not closed after every flush).
Shard names may include a worker prefix (w0-0) under parallel write; shuffled layouts add the bNNN- prefix.

Write-lane backpressure

Partition processors enqueue rows onto dataset write lanes (storage/run/write_lanes.py). When a lane backlog grows, the processor may log waiting for writer — that is intentional backpressure, not a crash. After the feed finishes, the run still drains lanes and finalizes parts/manifests.