# Checkpointing in Geneva

This document describes the **current checkpoint mechanism** used by Geneva’s Ray-based pipelines (backfill and materialized view refresh).

**Audience:** Contributors working on Ray pipeline execution, fault tolerance, and performance.

**See also:**
- `fragment_writers.md` (writer + commit flow)
- `materialized_views.md` (MV architecture, stable row IDs, refresh)

---

## Why checkpointing exists

Checkpointing is used to make long-running distributed jobs **idempotent and restartable**:

- **Fault tolerance**: if a worker/actor dies, we can resume without recomputing completed work.
- **Deduplication**: avoid redoing work across retries and incremental refreshes.
- **Incremental execution**: skip already-computed fragments/batches during re-runs.

At a high level, the pipeline produces **batch checkpoints** (RecordBatches) and the writer produces **fragment data-file checkpoints** (references to staged Lance files). A commit then atomically installs those files into the destination table.

### High-level diagram

![Geneva task graph](assets/geneva-task.png)

---

## Core abstractions

### CheckpointStore

All checkpoints are stored in a `CheckpointStore`, a key-value store of `pyarrow.RecordBatch`:

- **Interface**: `src/geneva/checkpoint.py::CheckpointStore`
- **Primary implementation**: `LanceCheckpointStore`, which stores each key as a `"{key}.lance"` file under a root directory (via `lance.file.LanceFileSession`).

Important constraint: **keys become file paths** in `LanceCheckpointStore`, so they must not contain dangerous path segments. Geneva avoids this by hashing user-provided strings (URI, where clause) inside key prefixes.

### ReadTask and MapTask

Checkpoint keys are generated by `MapTask` (not by `ReadTask`):

- `ReadTask` (`ScanTask`, `CopyTask`, …) defines *what to read* and its destination fragment/range:
  - `dest_frag_id()`
  - `dest_offset()`
  - `num_rows()`
- `MapTask` (`BackfillUDFTask`, `CopyTableTask`, …) defines *what to compute* and how to build checkpoint keys:
  - `checkpoint_prefix(...)`
  - `checkpoint_key(dataset_uri, dataset_version, frag_id, start, end, where)`

---

## Checkpoint key formats

Geneva uses **stable, structured keys** (see `src/geneva/checkpoint_utils.py`).

### Prefix

A prefix encodes the logical identity of work:

- **UDF/view name**
- **UDF/version** (changes when code changes)
- **Output column label**
- **where** hash (so different partial backfills don’t collide)
- **dataset URI** hash (so different tables don’t collide)
- **srcfiles hash** (hash of source data files for the input columns, when available)

Example prefix shape:

```
udf-{name}_ver-{version}_col-{column}_where-{md5(where)}_uri-{md5(uri)}_srcfiles-{md5(sorted input data files)}
```

Notes:
- The dataset version is intentionally *not* part of new checkpoint prefixes. Using
  dataset versions caused all checkpoints to invalidate after any write, even when
  inputs were unchanged.
- The srcfiles hash is derived from the **input columns'** source data files for
  the fragment. This tracks changes in the input data files without being tied to
  dataset versions.

### Per-batch checkpoint key (frag + range)

The canonical unit of checkpointed work is a **fragment-local range** `[start, end)`:

```
{prefix}_frag-{frag_id}_range-{start}-{end}
```

Key points:

- `frag_id` is the **destination** fragment ID.
- `start/end` are **fragment-local offsets** in the same domain as `ScanTask.offset/limit` (the domain used by `plan_read`).
- `end` is **exclusive**.

### Fragment data-file checkpoint key (dedupe key)

Once a fragment has been fully written, the pipeline records a “fragment is done” checkpoint that points at the staged `.lance` data file.

This key is computed by `src/geneva/runners/ray/pipeline.py::_get_fragment_dedupe_key(...)` and stores:

```python
{
  "file": new_file.path,
  "src_data_files": json.dumps(sorted(src_files)),       # MV refresh validation
  "output_field_ids": json.dumps(sorted(output_field_ids)),  # output column identity
}
```

This enables the next run to **skip** recomputing that fragment and still include it in the final commit.

For materialized views, the checkpoint may also store `src_data_files` (the union of
source data file paths for the inputs that contributed to the destination fragment).
This allows refresh logic to detect when source files change even if row IDs are
stable.

### Legacy keys

Older versions used different checkpoint naming. The current code keeps compatibility by:

- listing keys by prefix and parsing `_frag-{id}_range-{start}-{end}` when available
- falling back to legacy fragment keys (see `src/geneva/apply/__init__.py::_legacy_fragment_dedupe_key`)
- reconstructing per-batch checkpoints if a legacy “task-level” checkpoint exists

---

## How checkpoints are produced (applier stage)

### Execution

Ray `ApplierActor` runs a `CheckpointingApplier` over `ReadTask`s.

For each `ReadTask`, the applier:

1. Produces batches (with `_rowaddr` present for writers).
2. Applies the `MapTask` (UDFs / copy transforms).
3. Stores each resulting batch in the `CheckpointStore` under a **per-batch checkpoint key**.
4. Returns a `MapBatchCheckpoint` object containing:
   - `checkpoint_key`
   - `offset` (fragment-local)
   - `num_rows` (materialized rows in the stored batch)
   - `span` (how far the checkpoint advances in the planner’s offset domain)
   - `udf_rows` (rows actually processed by UDFs)

### Async checkpoint writer thread

UDF processing and checkpoint I/O run in **separate threads** so that UDF compute
is not blocked by object store writes. This also reduces request volume by
coalescing multiple small batches into fewer, larger writes.

#### Threading model

```mermaid
sequenceDiagram
    participant UDF as UDF Thread
    participant Q as queue.Queue
    participant W as Writer Thread

    UDF->>UDF: next(batch_iter)
    UDF->>UDF: _make_pending_batch_checkpoint()
    UDF->>Q: put(pending)
    Q->>W: get(timeout=flush_interval)
    W->>W: buffer pending item

    UDF->>UDF: next(batch_iter)
    UDF->>UDF: _make_pending_batch_checkpoint()
    UDF->>Q: put(pending)
    Q->>W: get(timeout=flush_interval)
    W->>W: buffer pending item

    Note over W: flush interval expires
    W->>W: _concat_record_batches()
    W->>W: _write_checkpoint_batch()

    UDF->>Q: put(SENTINEL)
    Q->>W: get() → SENTINEL
    W->>W: flush remaining & exit
    UDF->>W: join()
```

- The UDF thread produces `_PendingBatchCheckpoint` items and puts them on a
  `queue.Queue`.
- The writer thread (`_checkpoint_writer_loop`) consumes them, buffering items
  and flushing on a configurable time interval
  (`batch_checkpoint_flush_interval_seconds`, default 10 s).
- On flush, buffered batches are concatenated via `_concat_record_batches()` into
  a single `RecordBatch` before writing to the `CheckpointStore`, reducing the
  number of object store round-trips.
- A sentinel object (`_CHECKPOINT_FLUSH_SENTINEL`) signals the writer thread to
  flush remaining items and exit. The UDF thread always sends this in a `finally`
  block and joins the writer thread before returning.

#### Error handling

The writer thread captures exceptions into a shared `writer_errors` list. The UDF
thread checks this list after each batch and re-raises the first error as a
`RuntimeError` so failures surface promptly rather than being silently lost.

#### Adaptive sizing interaction

Because checkpoint I/O now happens off the UDF thread, wall-clock batch elapsed
time no longer reflects pure UDF cost. The adaptive checkpoint sizer
(`AdaptiveCheckpointSizer`) therefore uses **UDF-only elapsed time** — computed as
the delta of `batch_applier.udf_processing_time_ms` between batches — instead of
wall-clock time. This prevents slow object store writes from inflating the sizer's duration
estimate and shrinking checkpoint sizes unnecessarily.

### Range/span semantics (important for deletes and gaps)

A stored RecordBatch can have fewer rows than the checkpoint “coverage”. This happens when:

- the input has **deletes** (logical rows < physical rows), or
- the materialized output is **sparse** and the applier uses `_rowaddr` to indicate true physical placement.

To keep planning and ordering correct, the applier encodes the intended coverage in the checkpoint key’s `_range-{start}-{end}` and in `MapBatchCheckpoint.span`.

---

## Adaptive checkpoint sizing

Adaptive sizing lets the applier start with **small checkpoints** (faster proof‑of‑life and more frequent progress) and then grow checkpoint sizes as it observes runtime throughput, while never exceeding safe bounds.

### How it works

The applier wraps each `ReadTask` with `AdaptiveReadTask` and uses `AdaptiveCheckpointSizer`:

- The underlying read still fetches **max‑size batches** (bounded by `checkpoint_size` / `map_task.batch_size()`).
- Each max batch is **sliced** into smaller checkpoints based on the adaptive sizer’s current size.
- After each emitted checkpoint, the applier records the elapsed duration and row count, which updates the sizer for subsequent batches.

This means:
- **Planning** still uses the same `task_size` / `checkpoint_size` for estimating work.
- **Actual checkpoint spans** can be smaller than the planned `checkpoint_size` once adaptive sizing is enabled.

### Bounds and defaults

Adaptive sizing is always **clamped** to bounds:

- **`max_checkpoint_size`**: upper bound. Defaults to the job’s checkpoint size (`map_task.batch_size()`), and is also capped to that value if a larger max is provided.
- **`min_checkpoint_size`**: lower bound. Defaults to **1**.

When `min_checkpoint_size == max_checkpoint_size`, adaptive sizing is effectively disabled (fixed‑size checkpoints).

### Override points

You can override adaptive bounds at two levels:

1. **UDF definition** via `@udf(..., min_checkpoint_size=..., max_checkpoint_size=...)`
2. **Backfill call** via `table.backfill(..., min_checkpoint_size=..., max_checkpoint_size=...)`

Backfill‑level values take precedence over UDF defaults.

### Test stability

Many existing tests assume **one checkpoint per plan** (e.g., `len(results) == 1`). To keep those tests stable:

- Specify `min_checkpoint_size` and `max_checkpoint_size` to the same value (e.g., `DEFAULT_CHECKPOINT_ROWS`), or
- Set them on the test UDF / backfill call as needed.

---

## How checkpoints are used (planning stage)

Planning is done by `plan_read` (see `src/geneva/apply/__init__.py`). It:

1. Enumerates destination fragments.
2. Checks for a **fragment data-file checkpoint** (fragment-level dedupe key).
   - If present, it skips scheduling compute tasks for the fragment and collects an existing `DataFile` reference so the fragment can be included in commit.
3. Otherwise, it lists per-batch checkpoint keys, parses their `_range-` suffixes, merges covered ranges, and computes missing ranges.
   - Parsers match the *trailing* `_frag-{id}_range-{start}-{end}` suffix (splitting from the right where needed), so UDF/view names containing substrings like `_range-` do not cause conflicts.
4. Emits `ScanTask`s for only the missing ranges.

Result: retries and incremental refresh can be **O(missing work)** rather than O(total work).

### Output data-file validation (rebackfill safety)

When deciding whether a fragment-level checkpoint can be reused, we validate the
**output column identity** using field IDs stored in the checkpoint. This guards
against cases like `test_rebackfill`, where a column is dropped and re-added:

- The **input** source files (and thus `src_files_hash`) can remain unchanged.
- The **output** column data files change (new file UUIDs).

The src_files_hash tracks input data changes, but doesn't capture output schema
changes. When a column is dropped and re-added:

- The input src_files_hash is unchanged (same source data).
- The output column receives new field_ids from Lance schema evolution.
- The checkpointed data file was written with the old field_ids.
- Lance data files are tied to specific field_ids. A DataReplacementOperation cannot
  use a file written for field-id-X to populate a column expecting field-id-Y.

Note: Per-batch checkpoints (containing computed Arrow data) remain valid since
they're keyed by input hash, not output field_ids. A future optimization could
re-serialize these cached results to new files with correct field_ids, avoiding
UDF recomputation entirely.

Without validating output field IDs, we could incorrectly reuse a fragment-level
checkpoint and skip recomputation, leaving the re-added column as all `None`.

If a fragment-level checkpoint is present but the output field IDs no longer match,
the planner ignores any per-range checkpoints for that fragment and re-schedules work.

---

## How checkpoints are consumed (writer stage)

### Ordering and completeness

Fragment writers consume `(offset, checkpoint_key)` pairs through a Ray queue.

The writer must reassemble batches in **monotonic offset order**, even if applier results arrive out of order. This is done with `SequenceQueue`.

**Critical detail for deletes:**

- Writers expect offsets in the **logical-row domain** (`frag.count_rows()`), not `frag.physical_rows`.
- Physical layout is restored later using `_rowaddr`.

This is why the writer buffers until it has advanced through `num_logical_rows`, and then aligns to physical layout.

### Physical layout alignment

After reordering, the writer calls `_align_batches_to_physical_layout(...)` to ensure the output fragment has dense physical row addresses:

- intra-batch gaps are filled (`_fill_rowaddr_gaps`)
- inter-batch gaps are filled by inserting null rows for missing `_rowaddr` ranges

This is the step that makes materialized views with filters and backfills with deletes produce a valid physical fragment file.

### Column filtering

Before writing, batches are filtered to the destination schema columns (`_filter_columns_to_schema`) so writers don’t accidentally persist internal/extra columns.

---

## Commit and reuse

The pipeline collects `DataFile`s from fragment writers and commits them using:

- `lance.LanceOperation.DataReplacement`
- `lance.LanceDataset.commit(..., read_version=...)`

Fragments that were skipped due to fragment-level checkpoints are included as well, so the final table version is consistent.

---

## Failure modes and how checkpointing prevents them

### Worker/actor failures

- Applier failures: retries can load cached checkpoints and avoid recomputation.
- Writer failures: writers can be restarted; cached tasks and checkpoint keys can be replayed.

### Empty work units

If a `ReadTask` produces **no batches** (e.g., everything filtered out), the applier still persists “completion” checkpoints (null-filled batches) so the task is idempotent and the job can finish.

### Common hang symptom: logical vs physical row mismatch

If any stage mixes up:

- `frag.count_rows()` (logical rows, after deletes / filters)
- `frag.physical_rows` (physical rows)

…it can lead to a writer waiting forever for offsets that will never be produced. Writers should advance in the planner’s offset domain (logical rows) and only use `_rowaddr` to restore physical layout.

---

## Debugging tips

- **List keys**: `CheckpointStore.list_keys(prefix)` is the source of truth; keys are human-readable.
- **Look for `_range-` keys**: they indicate per-batch checkpoints.
- **Look for fragment dedupe keys**: they indicate a fragment has a staged datafile that can be reused.
- **If a job “hangs”**: inspect whether the writer is waiting for offsets beyond `frag.count_rows()` (logical rows) or whether checkpoint ranges have gaps.

---
