feat: add engine-native record selection

architecture/config.md

@@ -8,24 +8,29 @@ Source: `packages/data-designer-config/src/data_designer/config/`
 
 The config layer provides:
 - **`DataDesignerConfigBuilder`** — fluent builder for constructing dataset configs
-- **`DataDesignerConfig`** — the root config object holding columns, models, constraints, processors, and profilers
+- **`DataDesignerConfig`** — the root config object holding columns, models, constraints, processors, profilers, and an optional record-selection policy
 - **Column configs** — a discriminated union of Pydantic models, one per column type
 - **Model configs** — LLM endpoint configuration with inference parameters
 - **Sampler params** — statistical generator parameters with their own discriminated union
+- **Record selection** — an optional bounded policy for accepting generated candidate records
 - **Plugin integration** — runtime extension of config unions via entry-point plugins
 
 ## Key Components
 
 ### Builder API
 
-`DataDesignerConfigBuilder` is the primary construction surface. It holds mutable state (column configs, constraints, processors) and produces an immutable `DataDesignerConfig` on `build()`.
+`DataDesignerConfigBuilder` is the primary construction surface. It holds mutable state (column configs, constraints, processors, and record selection) and produces an immutable `DataDesignerConfig` on `build()`.
 
-- **Fluent mutators**: `add_column`, `add_constraint`, `add_processor`, `add_profiler`, `add_model_config`, `add_tool_config`, `with_seed_dataset`
+- **Fluent mutators**: `add_column`, `add_constraint`, `add_processor`, `add_profiler`, `add_model_config`, `add_tool_config`, `with_seed_dataset`, `with_record_selection`
 - **Column shorthand**: pass `name` + `column_type` + kwargs instead of a full config instance; the builder resolves the correct config class via `get_column_config_from_kwargs`
 - **Config loading**: `from_config` accepts dicts, file paths, URLs, or `BuilderConfig` objects; normalizes shorthand formats into the full structure
 
 `BuilderConfig` wraps `DataDesignerConfig` with a `library_version` field validated against the running version.
 
+### Record Selection
+
+`RecordSelectionConfig` declares an optional, bounded policy for accepting generated candidate records. It identifies a boolean predicate column, sets a strict positive `max_candidate_records` budget, and chooses the behavior when that budget is exhausted. The configuration is stored on `DataDesignerConfig` and attached through `DataDesignerConfigBuilder.with_record_selection()`; it remains declarative and does not call engine code.
+
 ### Column Configs
 
 All column configs inherit from `SingleColumnConfig(ConfigBase, ABC)`, which provides `name`, `drop`, `skip`, `propagate_skip`, and the `column_type` discriminator field.

architecture/dataset-builders.md

@@ -78,13 +78,15 @@ Manages per-row-group DataFrames and persistence:
 
 Checkpoint state lives in `metadata.json`. Each metadata write includes the config fingerprint (`config_hash`, `config_hash_algo`, and `config_hash_version`) so compatibility checks do not need to deserialize `builder_config.json` for the common path. `builder_config.json` remains the human-readable record of the run configuration and the fallback for older datasets.
 
-Resume scans `parquet-files/batch_*.parquet` and reads parquet metadata to recover the completed row-group IDs and their actual persisted row counts. `metadata.json` remains the source of truth for the run *configuration* (`buffer_size`, `target_num_records`, `original_target_num_records`, config fingerprint), but the filesystem is the source of truth for *progress* (`num_completed_batches`, `actual_num_records`). Splitting the two sources is what lets resume survive a crash between writing a row-group parquet and updating metadata - the filesystem reflects the durable state even when metadata lags by a step. Reading actual row counts also matters for early-shutdown salvage, where a completed parquet file can contain fewer rows than the requested row-group size. Resume tolerates non-contiguous IDs because row groups can complete out of order.
+For ordinary generation, resume scans `parquet-files/batch_*.parquet` and reads parquet metadata to recover the completed row-group IDs and their actual persisted row counts. `metadata.json` remains the source of truth for the run *configuration* (`buffer_size`, `target_num_records`, `original_target_num_records`, config fingerprint), but the filesystem is the source of truth for *progress* (`num_completed_batches`, `actual_num_records`). Splitting the two sources is what lets resume survive a crash between writing a row-group parquet and updating metadata - the filesystem reflects the durable state even when metadata lags by a step. Reading actual row counts also matters for early-shutdown salvage, where a completed parquet file can contain fewer rows than the requested row-group size. Resume tolerates non-contiguous IDs because row groups can complete out of order.
 
-Resume relies on stable row-group boundaries within a run. It treats datasets that have completed `process_after_generation()` as terminal: after-generation processors operate on the whole dataset and can re-chunk rows or change schema, invalidating row-group identity for later resume/extension. The terminal-state check raises a clear `DatasetGenerationError` (not a `TypeError`) when the persisted metadata is missing required fields such as `target_num_records`.
+Ordinary resume relies on stable row-group boundaries within a run. It treats datasets that have completed `process_after_generation()` as terminal: after-generation processors operate on the whole dataset and can re-chunk rows or change schema, invalidating row-group identity for later resume/extension. The terminal-state check raises a clear `DatasetGenerationError` (not a `TypeError`) when the persisted metadata is missing required fields such as `target_num_records`.
 
-After-generation processors run unconditionally on the on-disk dataset whenever they are configured — including the case where resume sees every row group already on disk. This closes the crash window between the final row-group parquet write and the `post_generation_state="started"` marker write: in that window, the dataset is complete but post-generation never ran, and the on-disk parquet files are still clean (no processor has touched them). The `post_generation_state="started"` short-circuit still rejects the other direction (`process_after_generation()` crashed mid-rewrite, leaving the parquet files in an ambiguous state), so resume only re-runs after-generation when it is safe to do so.
+On the ordinary path, after-generation processors run unconditionally on the on-disk dataset whenever they are configured — including the case where resume sees every row group already on disk. This closes the crash window between the final row-group parquet write and the `post_generation_state="started"` marker write: in that window, the dataset is complete but post-generation never ran, and the on-disk parquet files are still clean (no processor has touched them). The `post_generation_state="started"` short-circuit still rejects the other direction (`process_after_generation()` crashed mid-rewrite, leaving the parquet files in an ambiguous state), so resume only re-runs after-generation when it is safe to do so.
 
-Metadata writes are atomic (`tmp` file + `fsync` + `os.replace`) because `metadata.json` is the crash-recovery checkpoint. Corrupt or partially written metadata raises a clear `DatasetGenerationError` rather than falling through as a generic config mismatch.
+Record selection uses a separate two-phase checkpoint model. Each completed candidate batch commits a marker under `selection-checkpoints/` and, when it accepted rows, an immutable accepted-row partition under `selection-accepted/`; those artifacts, rather than the published `parquet-files/`, are the durable source of selection progress. Once enough rows are accepted—or the candidate budget is exhausted with `on_exhausted="return_partial"`—the engine publishes the final dataset from the immutable partitions and then runs after-generation processors. A valid completed publication is reused as-is on resume, while a missing or incomplete publication can be rebuilt from the immutable partitions without regenerating candidates. This keeps candidate progress resumable even when publication or after-generation processing is interrupted.
+
+Metadata writes are atomic (`tmp` file + `fsync` + `os.replace`). Metadata remains authoritative for run configuration and publication state; ordinary parquet files and record-selection markers are authoritative for generation progress. Corrupt or partially written metadata raises a clear `DatasetGenerationError` rather than falling through as a generic config mismatch.
 
 `DatasetCreationResults` from a resume invocation reflects the full on-disk dataset for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`), but per-run observability (`task_traces`, model-usage logs, telemetry events) is scoped to the current invocation — the original run's in-memory state is not persisted across process boundaries.
 

fern/versions/latest.yml

@@ -35,6 +35,8 @@ navigation:
         path: ./latest/pages/concepts/validators.mdx
       - page: Processors
         path: ./latest/pages/concepts/processors.mdx
+      - page: Record Selection
+        path: ./latest/pages/concepts/record-selection.mdx
       - page: Workflow Chaining
         path: ./latest/pages/concepts/workflow-chaining.mdx
       - page: Person Sampling

fern/versions/latest/pages/concepts/record-selection.mdx

@@ -0,0 +1,147 @@
+---
+title: "Record Selection"
+description: "Generate an exact number of rows that satisfy a declared boolean criterion."
+---
+
+# Record Selection
+
+Record selection lets one `DataDesigner.create()` call produce an exact number of rows that satisfy a declared
+boolean column. Data Designer owns candidate generation, filtering, refill, deterministic trimming, checkpoints,
+and resume.
+
+Use record selection when filtering after generation would leave fewer rows than you requested. For example, you
+can require a judge score, validation result, safety decision, or expression to pass before a row enters the output.
+
+## Configure an accepted-row target
+
+Declare the criterion as a normal boolean column, then attach `RecordSelectionConfig` to the builder:
+
+```python
+import data_designer.config as dd
+from data_designer.interface import DataDesigner
+
+builder = dd.DataDesignerConfigBuilder()
+builder.add_column(
+    dd.SamplerColumnConfig(
+        name="quality_score",
+        sampler_type=dd.SamplerType.UNIFORM,
+        params=dd.UniformSamplerParams(low=0.0, high=1.0),
+    )
+)
+builder.add_column(
+    dd.ExpressionColumnConfig(
+        name="meets_criteria",
+        expr="{{ quality_score >= 0.8 }}",
+        dtype="bool",
+        drop=True,
+    )
+)
+builder.with_record_selection(
+    dd.RecordSelectionConfig(
+        predicate_column="meets_criteria",
+        max_candidate_records=10_000,
+        on_exhausted="raise",
+    )
+)
+
+results = DataDesigner().create(builder, num_records=1_000)
+assert results.count_records() == 1_000
+```
+
+With record selection enabled, `num_records=1_000` means **1,000 accepted output rows**. Candidate generation stops
+after the completed batch that reaches that target, or when `max_candidate_records` is exhausted. Accepted rows keep
+candidate order; if the final candidate batch contains more passing rows than needed, Data Designer keeps the earliest
+rows.
+
+<Warning>
+  Candidate generation can invoke models for rows that are later rejected. Set `max_candidate_records` to a strict
+  positive integer that is greater than or equal to `num_records` and reflects your cost limit. Booleans, floats,
+  and numeric strings are rejected instead of being coerced into a budget.
+</Warning>
+
+## Predicate behavior
+
+The predicate column must exist in the dataset configuration. An expression predicate must use `dtype="bool"`.
+Category and subcategory sampler predicates are accepted only when every configured value, including conditional
+values, is boolean and no output conversion is configured. Other built-in non-boolean column types are rejected
+before generation. Seed, custom, and plugin columns with unknown output types are validated at runtime.
+
+| Predicate value | Result |
+| --- | --- |
+| `True` | Accept the row |
+| `False` | Reject the row |
+| Null | Reject the row and increment `null_predicate_records` |
+| Any other value | Stop with a generation error instead of applying truthiness coercion |
+
+Selection runs after the complete candidate-row DAG and before post-batch processors. A predicate with `drop=True`
+still participates in selection, then follows the normal dropped-column artifact policy.
+
+## Selection limitations
+
+Selection predicates must be row-local: they can depend on columns in the current candidate record, but not on the
+current batch, previously generated candidates, or previously accepted records. Global ranking, top-N selection,
+quotas, cross-batch deduplication, and stateful plugin predicates are not supported.
+
+Data Designer tracks and stages media produced by engine-managed image generators. Custom and plugin code should be
+side-effect-free: external files or media that it creates are not tracked by record selection, so artifacts associated
+with rejected candidates may remain.
+
+## Choose exhaustion behavior
+
+The default `on_exhausted="raise"` raises `DataDesignerRecordSelectionExhaustedError`. The error exposes the target,
+accepted count, candidate count, and configured cap:
+
+```python
+from data_designer.interface import DataDesignerRecordSelectionExhaustedError
+
+try:
+    results = DataDesigner().create(builder, num_records=1_000)
+except DataDesignerRecordSelectionExhaustedError as error:
+    print(error.accepted_records, error.candidate_records)
+```
+
+Use `on_exhausted="return_partial"` when a smaller accepted-only dataset is useful. A valid all-rejected result is
+a schema-bearing zero-row dataset. In that case, profiling is skipped and `results.load_analysis()` returns `None`.
+Early shutdown remains an error. A zero-row run with a non-retryable generation failure is also an error rather than
+an empty partial result; resume replays its durable terminal state.
+
+## Resume an interrupted selection run
+
+Candidate batches and accepted partitions are checkpointed independently, including candidate batches that accept
+zero rows. Resume continues from the next candidate offset without regenerating committed work:
+
+```python
+from data_designer.interface import ResumeMode
+
+results = DataDesigner().create(
+    builder,
+    num_records=1_000,
+    dataset_name="quality-filtered",
+    resume=ResumeMode.ALWAYS,
+)
+```
+
+For record selection, `ResumeMode.ALWAYS` requires the same configuration, `num_records`, and `RunConfig.buffer_size`
+as the original run. A larger target is not compatible because rows trimmed from the original final batch were not
+committed. `ResumeMode.IF_POSSIBLE` clears engine-managed artifacts and starts the selection run again when those
+inputs differ.
+
+Published output is rebuilt from immutable accepted partitions after a crash. `push_to_hub()` accepts only terminal
+selection artifacts whose publication state is complete; internal candidate markers and accepted-partition staging
+are not uploaded.
+
+<Note>
+  `preview()` does not execute the accepted-row retry and checkpoint contract. Preview the configuration before
+  enabling record selection, then use `create()` for the accepted-row target.
+</Note>
+
+## Processor limitations
+
+Pre-batch processors run on candidates. Record selection runs before post-batch processors, so post-batch processors
+see accepted rows only, must preserve their count, and must retain at least one column for non-empty output.
+After-generation processors run once on the published accepted dataset and have the same count and non-empty schema
+requirements. Data Designer validates both contracts and raises if a plugin violates them.
+
+Selection diagnostics are stored under `record_selection` in `metadata.json`, including candidate attempts, accepted,
+rejected, null, failed, and trimmed rows, acceptance rate, and terminal state. Model usage includes both accepted and
+rejected candidate work.

fern/versions/latest/pages/concepts/workflow-chaining.mdx

@@ -79,7 +79,7 @@ A stage can expose different views of its data:
 
 | Surface | What it returns |
 |---------|-----------------|
-| `results["stage_name"]` | The effective `DatasetCreationResults` for that stage. If the stage uses `output_processors`, this points at the output-processor run. |
+| `results["stage_name"]` | The effective `DatasetCreationResults` for that stage. For a non-empty stage with `output_processors`, this points at the output-processor run. An allowed empty stage skips output processors and retains the main stage result. |
 | `results.load_stage_output("stage_name")` | The selected output handed to downstream stages. This follows `output="processor:<name>"` and `on_success`. |
 | `results.load_dataset()` | The selected output from the final stage. |
 
@@ -136,7 +136,7 @@ workflow.add_stage(
 workflow.add_stage("enriched", enriched)
 ```
 
-`on_success_version` is part of the stage resume identity. Change it when the callback's output semantics change. If a callback returns zero rows, the workflow raises by default; set `allow_empty=True` to mark that stage as completed empty and skip downstream stages.
+`on_success_version` is part of the stage resume identity. Change it when the callback's output semantics change. If a selected stage output has zero rows, the workflow raises by default; set `allow_empty=True` to mark that stage as completed empty and skip downstream stages. When the main stage result is empty, configured output processors are not invoked; the schema-bearing main result is retained even when `output="processor:<name>"` was selected.
 
 ## Changing row counts between stages
 

packages/data-designer-config/src/data_designer/config/__init__.py

@@ -60,6 +60,10 @@
         ProcessorType,
         SchemaTransformProcessorConfig,
     )
+    from data_designer.config.record_selection import (  # noqa: F401
+        RecordSelectionConfig,
+        RecordSelectionExhaustion,
+    )
     from data_designer.config.run_config import (  # noqa: F401
         JinjaRenderingEngine,
         RequestAdmissionTuningConfig,
@@ -124,6 +128,7 @@
 _MOD_MCP = f"{_MOD_BASE}.mcp"
 _MOD_MODELS = f"{_MOD_BASE}.models"
 _MOD_PROCESSORS = f"{_MOD_BASE}.processors"
+_MOD_RECORD_SELECTION = f"{_MOD_BASE}.record_selection"
 _MOD_SAMPLER_CONSTRAINTS = f"{_MOD_BASE}.sampler_constraints"
 _MOD_SAMPLER_PARAMS = f"{_MOD_BASE}.sampler_params"
 _MOD_SEED = f"{_MOD_BASE}.seed"
@@ -187,6 +192,9 @@
     "DropColumnsProcessorConfig": (_MOD_PROCESSORS, "DropColumnsProcessorConfig"),
     "ProcessorType": (_MOD_PROCESSORS, "ProcessorType"),
     "SchemaTransformProcessorConfig": (_MOD_PROCESSORS, "SchemaTransformProcessorConfig"),
+    # record selection
+    "RecordSelectionConfig": (_MOD_RECORD_SELECTION, "RecordSelectionConfig"),
+    "RecordSelectionExhaustion": (_MOD_RECORD_SELECTION, "RecordSelectionExhaustion"),
     # run_config
     "JinjaRenderingEngine": (f"{_MOD_BASE}.run_config", "JinjaRenderingEngine"),
     "RequestAdmissionTuningConfig": (f"{_MOD_BASE}.run_config", "RequestAdmissionTuningConfig"),