Add keyed execution for group-scoped consistent generation

[andreatnvidia]

Priority Level

Medium

Task Summary

Add a first-class keyed-execution policy that lets supported column generators run once per logical group and broadcast the generated value and side-effect columns to every row with the same key.

The policy should compose with existing column configurations so users can apply group-scoped consistency to built-in generators without creating a new wrapper implementation for each column type.

Acceptance criteria:

• Existing unscoped column behavior remains unchanged.
• A column can declare one or more upstream key columns.
• Concurrent rows with the same key invoke the underlying generator once.
• Results remain consistent across row groups and resumed runs.
• Primary outputs and declared side-effect columns are broadcast together.
• Conflicting dependency values within one key produce a clear error by default.
• Existing generator scheduling metadata, retries, validation, and model-provider behavior remain in effect.

Example

Suppose a seed dataset contains multiple events for the same account:

account_id	first_name	last_name	event
acct-1	Amina	Diallo	signup
acct-2	Carlos	Silva	signup
acct-1	Amina	Diallo	purchase

The user wants to reuse the existing LLM column generator while producing one consistent contact email per account. The API names below are illustrative:

builder.add_column(
    LLMTextColumnConfig(
        name="contact_email",
        model_alias="email-generator",
        prompt=(
            "Generate a plausible email for "
            "{{ first_name }} {{ last_name }}."
        ),
        consistency=KeyedExecutionConfig(
            key_columns=["account_id"],
            conflict_policy="error",
            persist=True,
        ),
    ),
)

Expected behavior:

• The underlying LLM generator is called twice, once for each unique account_id.
• Both acct-1 rows receive the same generated email.
• The mapping remains available when generation resumes.
• The existing LLM generator retains its provider admission, retries, validation, and trace behavior.
• If two acct-1 rows contain different first_name or last_name values, generation fails with a dependency-conflict error instead of silently choosing one.

The important requirement is that scoping wraps the existing column configuration rather than introducing a separate scoped implementation for each generator type.

Technical Details & Implementation Plan

1. Add an optional keyed-execution configuration to supported single-column configs. The initial shape should include key columns, an optional namespace, persistence behavior, and a dependency-conflict policy.
2. Teach the execution graph to treat key columns as dependencies so a key is computed only after its inputs are available.
3. Add a scheduler-owned keyed result store and single-flight coordination. The first ready row for a missing key invokes the existing generator; concurrent followers await and reuse that result.
4. Store the complete generated output for the column, including side-effect columns and an input-dependency fingerprint.
5. Persist keyed mappings through artifact storage and reload them during resume. Scope the mapping identity by configuration fingerprint, column, namespace, and canonicalized key.
6. Default dependency conflicts to error. A later design may support explicit alternatives such as selecting a representative row or incorporating dependencies into the effective key.
7. Implement phase one for cell-by-cell single-column generators. This covers LLM text, structured, code and judge columns, image and embedding columns, and compatible custom generators.
8. Evaluate full-column generators separately. Multi-column samplers, seed generation, and other from-scratch generators should remain out of scope until their semantics are defined.

Tests should cover non-contiguous rows, groups spanning row groups, concurrent duplicate keys, retry behavior, process resume, configuration changes, side-effect columns, skips, null keys, composite keys, and dependency conflicts.

Investigation / Context

Data Designer currently schedules cell-by-cell generators once per row and full-column generators once per row group. The dataset builder resolves each config type to one generator implementation, while plugins register additional config and implementation pairs.

A plugin can therefore add a specialized group-aware generator, but it cannot transparently apply group scoping to existing built-in generators while preserving their execution strategy and scheduler behavior. Implementing consistency as a core execution policy avoids per-generator wrappers and keeps provider admission, retries, dependency ordering, and side effects owned by Data Designer.

Agent Plan / Findings

A staged implementation is recommended:

• Phase one: cell-by-cell single-column generators with in-memory single-flight reuse, durable mappings, strict conflict detection, and resume coverage.
• Phase two: optimize scheduler admission so cache followers do not consume model-request capacity unnecessarily.
• Phase three: define and add support for compatible full-column generators.

The existing per-row task model can be preserved initially. Key lookup and single-flight coordination can wrap the generator invocation path, with each row task still completing independently after the shared result is applied.

Dependencies

None.