diff --git a/architecture/dataset-builders.md b/architecture/dataset-builders.md index 9c10cf531..0cca1c7d8 100644 --- a/architecture/dataset-builders.md +++ b/architecture/dataset-builders.md @@ -40,10 +40,11 @@ Preparation (`_prepare_async_run`): ### Execution Graph `ExecutionGraph` (in `dataset_builders/utils/execution_graph.py`) models column dependencies: -- Upstream/downstream sets derived from `required_columns` and side-effect columns +- Upstream/downstream sets derived from `required_columns`, side-effect columns, and `skip.when` references - `GenerationStrategy` per column (CELL_BY_CELL or FULL_COLUMN) - Kahn topological sort for execution order - `split_upstream_by_strategy` — separates batch-level from cell-level dependencies +- Skip metadata per column — `get_skip_config`, `should_propagate_skip`, `get_required_columns`, and `get_side_effect_columns` — queried at runtime by both engines to evaluate skip decisions ### CompletionTracker @@ -53,9 +54,28 @@ Tracks per-row-group, per-column completion state: - **Frontier**: computes ready tasks when backed by `ExecutionGraph` - Handles dropped rows and downstream task enqueuing +### Conditional Generation (Skip) + +Columns can be conditionally skipped per-row via `SkipConfig` (defined in `data_designer.config.base`). Two mechanisms control skipping: + +1. **Expression gate** — `skip=SkipConfig(when="{{ expr }}")` on a `SingleColumnConfig`. The Jinja2 expression is evaluated per-row; when truthy, the column is skipped for that row and the configured `value` (default `None`) is written instead of calling the generator. +2. **Skip propagation** — when an upstream column was skipped, downstream columns auto-skip unless they set `propagate_skip=False`. Propagation checks `required_columns` against the row's `__internal_skipped_columns` set. + +Skip evaluation is handled by two utility modules: + +- **`skip_evaluator.py`** — `evaluate_skip_when` renders the expression in a `NativeSandboxedEnvironment` (native Python types, `StrictUndefined`). `should_skip_by_propagation` checks set intersection between required columns and skipped columns. +- **`skip_tracker.py`** — manages the `__internal_skipped_columns` metadata key on record dicts. Each record carries a `__internal_skipped_columns` set listing which columns were skipped for that row. `apply_skip_to_record` adds the column name to that set, writes the skip value into the cell, and clears any side-effect columns. `strip_skip_metadata_from_records` removes the `__internal_skipped_columns` key before DataFrame construction so it never reaches parquet (called by `DatasetBatchManager`, `RowGroupBufferManager`, and inline in both engines). + +Both execution modes integrate skip at the same points: + +- **Sequential**: `_run_full_column_generator` and the fan-out methods (`_fan_out_with_threads`, `_fan_out_with_async`) call `_should_skip_cell` per record. Skipped rows are excluded from the generator input, then merged back with skip metadata preserved. A fast `_column_can_skip` check short-circuits the per-record evaluation when no skip config or propagation applies. +- **Async**: `_run_cell` and `_run_batch` in `AsyncTaskScheduler` call `_should_skip_record` / `_apply_skip_to_record` with the same logic. Skipped cells report as skipped (not success) in progress tracking. + +DAG edges are added for `skip.when` column references (both in `dag.py` and `ExecutionGraph.create`) so skip-gate columns are generated before the gated column. + ### DAG (Config-Level) -`dataset_builders/utils/dag.py` provides `topologically_sort_column_configs` — builds a NetworkX graph from `required_columns` and side-effect columns, returns a topological ordering. Used by both execution modes for initial column ordering. +`dataset_builders/utils/dag.py` provides `topologically_sort_column_configs` — builds a NetworkX graph from `required_columns`, side-effect columns, and `skip.when` references, returns a topological ordering. Used by both execution modes for initial column ordering. ### DatasetBatchManager diff --git a/docs/colab_notebooks/1-the-basics.ipynb b/docs/colab_notebooks/1-the-basics.ipynb index 14694d011..056a3c7db 100644 --- a/docs/colab_notebooks/1-the-basics.ipynb +++ b/docs/colab_notebooks/1-the-basics.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "527e4e8f", + "id": "9a82d43a", "metadata": {}, "source": [ "\"Open" @@ -10,7 +10,7 @@ }, { "cell_type": "markdown", - "id": "e58e7a85", + "id": "4d30e1c7", "metadata": {}, "source": [ "# 🎨 Data Designer Tutorial: The Basics\n", @@ -22,7 +22,7 @@ }, { "cell_type": "markdown", - "id": "199610f6", + "id": "f2a53a3c", "metadata": {}, "source": [ "### 📦 Import Data Designer\n", @@ -34,7 +34,7 @@ }, { "cell_type": "markdown", - "id": "9b738a91", + "id": "c442284e", "metadata": {}, "source": [ "### ⚡ Colab Setup\n", @@ -45,7 +45,7 @@ { "cell_type": "code", "execution_count": null, - "id": "af2bf6f9", + "id": "dac0f01a", "metadata": {}, "outputs": [], "source": [ @@ -56,7 +56,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0640f8fd", + "id": "5da0d2a0", "metadata": {}, "outputs": [], "source": [ @@ -74,7 +74,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e84edbc3", + "id": "d20ed2cb", "metadata": {}, "outputs": [], "source": [ @@ -84,7 +84,7 @@ }, { "cell_type": "markdown", - "id": "32d99706", + "id": "fdb97cba", "metadata": {}, "source": [ "### ⚙️ Initialize the Data Designer interface\n", @@ -97,7 +97,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e998f3ae", + "id": "fc8f76cf", "metadata": {}, "outputs": [], "source": [ @@ -106,7 +106,7 @@ }, { "cell_type": "markdown", - "id": "6afad397", + "id": "67bc996a", "metadata": {}, "source": [ "### 🎛️ Define model configurations\n", @@ -123,7 +123,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14ab53fd", + "id": "70d5cddd", "metadata": {}, "outputs": [], "source": [ @@ -153,7 +153,7 @@ }, { "cell_type": "markdown", - "id": "a73fef78", + "id": "e8e02b51", "metadata": {}, "source": [ "### 🏗️ Initialize the Data Designer Config Builder\n", @@ -168,7 +168,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5367d8c8", + "id": "2fdd1312", "metadata": {}, "outputs": [], "source": [ @@ -177,7 +177,7 @@ }, { "cell_type": "markdown", - "id": "c9bec24b", + "id": "4056cfd3", "metadata": {}, "source": [ "## 🎲 Getting started with sampler columns\n", @@ -194,7 +194,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b5551231", + "id": "85dd5c70", "metadata": {}, "outputs": [], "source": [ @@ -203,7 +203,7 @@ }, { "cell_type": "markdown", - "id": "9076cef7", + "id": "541b2033", "metadata": {}, "source": [ "Let's start designing our product review dataset by adding product category and subcategory columns.\n" @@ -212,7 +212,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5fdddb75", + "id": "356b361e", "metadata": {}, "outputs": [], "source": [ @@ -293,7 +293,7 @@ }, { "cell_type": "markdown", - "id": "282a37f8", + "id": "287cf22a", "metadata": {}, "source": [ "Next, let's add samplers to generate data related to the customer and their review.\n" @@ -302,7 +302,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5d018574", + "id": "282d9074", "metadata": {}, "outputs": [], "source": [ @@ -339,7 +339,7 @@ }, { "cell_type": "markdown", - "id": "5d4dee03", + "id": "102c1634", "metadata": {}, "source": [ "## 🦜 LLM-generated columns\n", @@ -354,7 +354,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f6e79f81", + "id": "20dc1332", "metadata": {}, "outputs": [], "source": [ @@ -390,7 +390,7 @@ }, { "cell_type": "markdown", - "id": "f562005f", + "id": "a983bf8a", "metadata": {}, "source": [ "### 🔁 Iteration is key – preview the dataset!\n", @@ -407,7 +407,7 @@ { "cell_type": "code", "execution_count": null, - "id": "70d761cd", + "id": "6b019b6e", "metadata": {}, "outputs": [], "source": [ @@ -417,7 +417,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6b1c75a5", + "id": "82ab36be", "metadata": {}, "outputs": [], "source": [ @@ -428,7 +428,7 @@ { "cell_type": "code", "execution_count": null, - "id": "77d0530c", + "id": "a75256be", "metadata": {}, "outputs": [], "source": [ @@ -438,7 +438,7 @@ }, { "cell_type": "markdown", - "id": "9c22fe3a", + "id": "d6d92058", "metadata": {}, "source": [ "### 📊 Analyze the generated data\n", @@ -451,7 +451,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8619efbb", + "id": "63b19fbc", "metadata": {}, "outputs": [], "source": [ @@ -461,7 +461,7 @@ }, { "cell_type": "markdown", - "id": "7d538cfd", + "id": "66073eb7", "metadata": {}, "source": [ "### 🆙 Scale up!\n", @@ -474,7 +474,7 @@ { "cell_type": "code", "execution_count": null, - "id": "af3702b8", + "id": "9270d1fc", "metadata": {}, "outputs": [], "source": [ @@ -484,7 +484,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c862c183", + "id": "77b7dec0", "metadata": {}, "outputs": [], "source": [ @@ -497,7 +497,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5228e949", + "id": "831eed73", "metadata": {}, "outputs": [], "source": [ @@ -509,14 +509,14 @@ }, { "cell_type": "markdown", - "id": "1d434fd7", + "id": "54a22faf", "metadata": {}, "source": [ "## ⏭️ Next Steps\n", "\n", "Now that you've seen the basics of Data Designer, check out the following notebooks to learn more about:\n", "\n", - "- [Structured outputs and jinja expressions](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/2-structured-outputs-and-jinja-expressions/)\n", + "- [Structured outputs, jinja expressions, and conditional generation](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/2-structured-outputs-and-jinja-expressions/)\n", "\n", "- [Seeding synthetic data generation with an external dataset](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/3-seeding-with-a-dataset/)\n", "\n", diff --git a/docs/colab_notebooks/2-structured-outputs-and-jinja-expressions.ipynb b/docs/colab_notebooks/2-structured-outputs-and-jinja-expressions.ipynb index eb166ee70..194ace96d 100644 --- a/docs/colab_notebooks/2-structured-outputs-and-jinja-expressions.ipynb +++ b/docs/colab_notebooks/2-structured-outputs-and-jinja-expressions.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "1ad92c47", + "id": "2304ee0b", "metadata": {}, "source": [ "\"Open" @@ -10,21 +10,21 @@ }, { "cell_type": "markdown", - "id": "940c1155", + "id": "4a4bd6b8", "metadata": {}, "source": [ - "# 🎨 Data Designer Tutorial: Structured Outputs and Jinja Expressions\n", + "# 🎨 Data Designer Tutorial: Structured Outputs, Jinja Expressions, and Conditional Generation\n", "\n", "#### 📚 What you'll learn\n", "\n", - "In this notebook, we will continue our exploration of Data Designer, demonstrating more advanced data generation using structured outputs and Jinja expressions.\n", + "In this notebook, we will continue our exploration of Data Designer, demonstrating more advanced data generation using structured outputs, Jinja expressions, and conditional generation with `skip.when`.\n", "\n", "If this is your first time using Data Designer, we recommend starting with the [first notebook](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/1-the-basics/) in this tutorial series.\n" ] }, { "cell_type": "markdown", - "id": "24bf6353", + "id": "a538f55f", "metadata": {}, "source": [ "### 📦 Import Data Designer\n", @@ -36,7 +36,7 @@ }, { "cell_type": "markdown", - "id": "4acc5938", + "id": "5a84258e", "metadata": {}, "source": [ "### ⚡ Colab Setup\n", @@ -47,7 +47,7 @@ { "cell_type": "code", "execution_count": null, - "id": "158affe9", + "id": "e92fca3c", "metadata": {}, "outputs": [], "source": [ @@ -58,7 +58,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a957fda7", + "id": "a50ce276", "metadata": {}, "outputs": [], "source": [ @@ -76,7 +76,7 @@ { "cell_type": "code", "execution_count": null, - "id": "06ba3a03", + "id": "72f7c8b5", "metadata": {}, "outputs": [], "source": [ @@ -86,7 +86,7 @@ }, { "cell_type": "markdown", - "id": "7cda7d10", + "id": "eb148ee6", "metadata": {}, "source": [ "### ⚙️ Initialize the Data Designer interface\n", @@ -99,7 +99,7 @@ { "cell_type": "code", "execution_count": null, - "id": "edfc7810", + "id": "3ee03591", "metadata": {}, "outputs": [], "source": [ @@ -108,7 +108,7 @@ }, { "cell_type": "markdown", - "id": "94a465c9", + "id": "e88a8757", "metadata": {}, "source": [ "### 🎛️ Define model configurations\n", @@ -125,7 +125,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f251f24b", + "id": "a8328629", "metadata": {}, "outputs": [], "source": [ @@ -155,7 +155,7 @@ }, { "cell_type": "markdown", - "id": "3492d64d", + "id": "7345cd69", "metadata": {}, "source": [ "### 🏗️ Initialize the Data Designer Config Builder\n", @@ -170,7 +170,7 @@ { "cell_type": "code", "execution_count": null, - "id": "91d0af65", + "id": "81827e18", "metadata": {}, "outputs": [], "source": [ @@ -179,7 +179,7 @@ }, { "cell_type": "markdown", - "id": "c30bd4b1", + "id": "018d1fd0", "metadata": {}, "source": [ "### 🧑‍🎨 Designing our data\n", @@ -206,7 +206,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a576dd08", + "id": "46643d0b", "metadata": {}, "outputs": [], "source": [ @@ -234,7 +234,7 @@ }, { "cell_type": "markdown", - "id": "a13cb8b4", + "id": "32438a65", "metadata": {}, "source": [ "Next, let's design our product review dataset using a few more tricks compared to the previous notebook.\n" @@ -243,7 +243,7 @@ { "cell_type": "code", "execution_count": null, - "id": "847185eb", + "id": "fbe1d47d", "metadata": {}, "outputs": [], "source": [ @@ -352,7 +352,7 @@ }, { "cell_type": "markdown", - "id": "974f8bf6", + "id": "0945f1a5", "metadata": {}, "source": [ "Next, we will use more advanced Jinja expressions to create new columns.\n", @@ -369,7 +369,7 @@ { "cell_type": "code", "execution_count": null, - "id": "370621ad", + "id": "76945654", "metadata": {}, "outputs": [], "source": [ @@ -422,7 +422,130 @@ }, { "cell_type": "markdown", - "id": "c0b4600b", + "id": "d71c91c3", + "metadata": {}, + "source": [ + "## 🚦 Conditional generation with `skip.when`\n", + "\n", + "So far, every column is generated for every row. But sometimes an expensive LLM column only makes sense\n", + "for a subset of rows — for example, a detailed complaint analysis is only useful when the review is negative.\n", + "\n", + "Data Designer lets you **skip** column generation on a per-row basis using `SkipConfig`.\n", + "Skipped rows receive `None` by default, but you can provide a sentinel value with\n", + "`skip=dd.SkipConfig(when=\"...\", value=\"N/A\")` to write a specific value instead.\n", + "\n", + "There are three patterns to know:\n", + "\n", + "| Pattern | How | Effect |\n", + "|---|---|---|\n", + "| **Expression gate** | `skip=dd.SkipConfig(when=\"...\")` | Skip this column when the Jinja2 expression is truthy |\n", + "| **Skip propagation** (default) | Downstream column depends on a skipped column | Automatically skipped too (`propagate_skip=True` by default) |\n", + "| **Propagation opt-out** | `propagate_skip=False` on the downstream column | Always generates, even if an upstream was skipped |\n" + ] + }, + { + "cell_type": "markdown", + "id": "0fed3486", + "metadata": {}, + "source": [ + "**Pattern 1 — Expression gate.** Only generate a detailed complaint analysis when the customer gave a low rating (1 or 2 stars).\n", + "Rows where the rating is 3 or higher will get `None` for this column.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d94412c1", + "metadata": {}, + "outputs": [], + "source": [ + "config_builder.add_column(\n", + " dd.LLMTextColumnConfig(\n", + " name=\"complaint_analysis\",\n", + " model_alias=MODEL_ALIAS,\n", + " prompt=(\n", + " \"A customer reviewed '{{ product.name }}' ({{ product_category }} / {{ product_subcategory }}).\\n\\n\"\n", + " \"Review: {{ customer_review.review }}\\n\"\n", + " \"Rating: {{ customer_review.rating }}/5\\n\"\n", + " \"Mood: {{ customer_review.customer_mood }}\\n\\n\"\n", + " \"Write a short root-cause analysis of why this customer is unhappy \"\n", + " \"and suggest one concrete improvement the product team could make.\"\n", + " ),\n", + " skip=dd.SkipConfig(when=\"{{ customer_review.rating > 2 }}\"),\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "61105564", + "metadata": {}, + "source": [ + "**Pattern 2 — Skip propagation.** `action_items` depends on `complaint_analysis`.\n", + "When `complaint_analysis` is skipped, `action_items` auto-skips too because\n", + "`propagate_skip` defaults to `True`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "07325bb5", + "metadata": {}, + "outputs": [], + "source": [ + "config_builder.add_column(\n", + " dd.LLMTextColumnConfig(\n", + " name=\"action_items\",\n", + " model_alias=MODEL_ALIAS,\n", + " prompt=(\n", + " \"Based on this complaint analysis:\\n\"\n", + " \"{{ complaint_analysis }}\\n\\n\"\n", + " \"List 2-3 concrete action items for the product team.\"\n", + " ),\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "0ba75f20", + "metadata": {}, + "source": [ + "**Pattern 3 — Propagation opt-out.** `review_summary` also depends on `complaint_analysis`,\n", + "but sets `propagate_skip=False` so it always generates. The prompt uses a Jinja conditional\n", + "to handle the case where `complaint_analysis` is `None`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36aa8ff1", + "metadata": {}, + "outputs": [], + "source": [ + "config_builder.add_column(\n", + " dd.LLMTextColumnConfig(\n", + " name=\"review_summary\",\n", + " model_alias=MODEL_ALIAS,\n", + " propagate_skip=False,\n", + " prompt=(\n", + " \"Summarize this product review in one sentence:\\n\"\n", + " \"Product: {{ product.name }}\\n\"\n", + " \"Rating: {{ customer_review.rating }}/5\\n\"\n", + " \"Review: {{ customer_review.review }}\\n\"\n", + " \"{% if complaint_analysis %}\"\n", + " \"Complaint analysis: {{ complaint_analysis }}\\n\"\n", + " \"{% endif %}\"\n", + " ),\n", + " )\n", + ")\n", + "\n", + "data_designer.validate(config_builder)" + ] + }, + { + "cell_type": "markdown", + "id": "0f65bf18", "metadata": {}, "source": [ "### 🔁 Iteration is key – preview the dataset!\n", @@ -439,7 +562,7 @@ { "cell_type": "code", "execution_count": null, - "id": "74ba0f8d", + "id": "6e8ce368", "metadata": {}, "outputs": [], "source": [ @@ -449,28 +572,32 @@ { "cell_type": "code", "execution_count": null, - "id": "70e7b7d1", + "id": "bf6a10a4", "metadata": {}, "outputs": [], "source": [ "# Run this cell multiple times to cycle through the 2 preview records.\n", + "# Look for rows where complaint_analysis and action_items are None (skipped)\n", + "# vs rows where they were generated (low-rated reviews).\n", "preview.display_sample_record()" ] }, { "cell_type": "code", "execution_count": null, - "id": "c9ca2e13", + "id": "00657954", "metadata": {}, "outputs": [], "source": [ "# The preview dataset is available as a pandas DataFrame.\n", + "# Notice that complaint_analysis, action_items, and review_summary columns\n", + "# reflect the skip behavior: None for skipped rows, generated text otherwise.\n", "preview.dataset" ] }, { "cell_type": "markdown", - "id": "c2f322ef", + "id": "b1f67cfe", "metadata": {}, "source": [ "### 📊 Analyze the generated data\n", @@ -483,7 +610,7 @@ { "cell_type": "code", "execution_count": null, - "id": "cbd8af52", + "id": "4d8d4eec", "metadata": {}, "outputs": [], "source": [ @@ -493,7 +620,7 @@ }, { "cell_type": "markdown", - "id": "1b0f1d16", + "id": "35bf8910", "metadata": {}, "source": [ "### 🆙 Scale up!\n", @@ -506,7 +633,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18cdcf71", + "id": "6a251f72", "metadata": {}, "outputs": [], "source": [ @@ -516,7 +643,7 @@ { "cell_type": "code", "execution_count": null, - "id": "bae6816f", + "id": "5c53c7ee", "metadata": {}, "outputs": [], "source": [ @@ -529,7 +656,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2ba65001", + "id": "806a6fde", "metadata": {}, "outputs": [], "source": [ @@ -541,7 +668,7 @@ }, { "cell_type": "markdown", - "id": "9c72488e", + "id": "c4d9093c", "metadata": {}, "source": [ "## ⏭️ Next Steps\n", diff --git a/docs/colab_notebooks/3-seeding-with-a-dataset.ipynb b/docs/colab_notebooks/3-seeding-with-a-dataset.ipynb index eae9687b1..a1d74c9cb 100644 --- a/docs/colab_notebooks/3-seeding-with-a-dataset.ipynb +++ b/docs/colab_notebooks/3-seeding-with-a-dataset.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "5bac5b4b", + "id": "38b44faa", "metadata": {}, "source": [ "\"Open" @@ -10,7 +10,7 @@ }, { "cell_type": "markdown", - "id": "fb608ca3", + "id": "be41c209", "metadata": {}, "source": [ "# 🎨 Data Designer Tutorial: Seeding Synthetic Data Generation with an External Dataset\n", @@ -24,7 +24,7 @@ }, { "cell_type": "markdown", - "id": "fefad44f", + "id": "858cfcd2", "metadata": {}, "source": [ "### 📦 Import Data Designer\n", @@ -36,7 +36,7 @@ }, { "cell_type": "markdown", - "id": "5187c796", + "id": "f422bbe3", "metadata": {}, "source": [ "### ⚡ Colab Setup\n", @@ -47,7 +47,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5b412c3f", + "id": "ad310340", "metadata": {}, "outputs": [], "source": [ @@ -58,7 +58,7 @@ { "cell_type": "code", "execution_count": null, - "id": "452e231a", + "id": "94b46509", "metadata": {}, "outputs": [], "source": [ @@ -76,7 +76,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f2c57532", + "id": "6acf8163", "metadata": {}, "outputs": [], "source": [ @@ -86,7 +86,7 @@ }, { "cell_type": "markdown", - "id": "5f39d576", + "id": "7b1f0367", "metadata": {}, "source": [ "### ⚙️ Initialize the Data Designer interface\n", @@ -99,7 +99,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8c3facea", + "id": "40156eb6", "metadata": {}, "outputs": [], "source": [ @@ -108,7 +108,7 @@ }, { "cell_type": "markdown", - "id": "a4371f10", + "id": "fe4b26e2", "metadata": {}, "source": [ "### 🎛️ Define model configurations\n", @@ -125,7 +125,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f6c50e67", + "id": "6b80e566", "metadata": {}, "outputs": [], "source": [ @@ -155,7 +155,7 @@ }, { "cell_type": "markdown", - "id": "1879ce5e", + "id": "1c4a6421", "metadata": {}, "source": [ "### 🏗️ Initialize the Data Designer Config Builder\n", @@ -170,7 +170,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7597e069", + "id": "eeef350c", "metadata": {}, "outputs": [], "source": [ @@ -179,7 +179,7 @@ }, { "cell_type": "markdown", - "id": "20146510", + "id": "862e3bb6", "metadata": {}, "source": [ "## 🏥 Prepare a seed dataset\n", @@ -204,7 +204,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ede54203", + "id": "2107b15c", "metadata": {}, "outputs": [], "source": [ @@ -222,7 +222,7 @@ }, { "cell_type": "markdown", - "id": "592bcb87", + "id": "e5ee5000", "metadata": {}, "source": [ "## 🎨 Designing our synthetic patient notes dataset\n", @@ -235,7 +235,7 @@ { "cell_type": "code", "execution_count": null, - "id": "91758b6a", + "id": "18e9c7ea", "metadata": {}, "outputs": [], "source": [ @@ -316,7 +316,7 @@ }, { "cell_type": "markdown", - "id": "7f05ab3c", + "id": "3b12313b", "metadata": {}, "source": [ "### 🔁 Iteration is key – preview the dataset!\n", @@ -333,7 +333,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d3fcc610", + "id": "133c96c1", "metadata": {}, "outputs": [], "source": [ @@ -343,7 +343,7 @@ { "cell_type": "code", "execution_count": null, - "id": "53a774ad", + "id": "4754788f", "metadata": {}, "outputs": [], "source": [ @@ -354,7 +354,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f61284e9", + "id": "3f2f1710", "metadata": {}, "outputs": [], "source": [ @@ -364,7 +364,7 @@ }, { "cell_type": "markdown", - "id": "eadc86f1", + "id": "b6534304", "metadata": {}, "source": [ "### 📊 Analyze the generated data\n", @@ -377,7 +377,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3dc714a5", + "id": "5aac62d0", "metadata": {}, "outputs": [], "source": [ @@ -387,7 +387,7 @@ }, { "cell_type": "markdown", - "id": "48fa3db8", + "id": "43c5d0a2", "metadata": {}, "source": [ "### 🆙 Scale up!\n", @@ -400,7 +400,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a8426985", + "id": "6324561c", "metadata": {}, "outputs": [], "source": [ @@ -410,7 +410,7 @@ { "cell_type": "code", "execution_count": null, - "id": "64bc95ab", + "id": "178f0d22", "metadata": {}, "outputs": [], "source": [ @@ -423,7 +423,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e5882e6d", + "id": "625ea810", "metadata": {}, "outputs": [], "source": [ @@ -435,7 +435,7 @@ }, { "cell_type": "markdown", - "id": "26b5a23f", + "id": "2dbe3006", "metadata": {}, "source": [ "## ⏭️ Next Steps\n", diff --git a/docs/colab_notebooks/4-providing-images-as-context.ipynb b/docs/colab_notebooks/4-providing-images-as-context.ipynb index 1ab12a837..be603eb7d 100644 --- a/docs/colab_notebooks/4-providing-images-as-context.ipynb +++ b/docs/colab_notebooks/4-providing-images-as-context.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "276ed0fd", + "id": "be521ec5", "metadata": {}, "source": [ "\"Open" @@ -10,7 +10,7 @@ }, { "cell_type": "markdown", - "id": "c4c71d59", + "id": "9af469a2", "metadata": {}, "source": [ "# 🎨 Data Designer Tutorial: Providing Images as Context for Vision-Based Data Generation" @@ -18,7 +18,7 @@ }, { "cell_type": "markdown", - "id": "cc3d5d34", + "id": "4ecf30a2", "metadata": {}, "source": [ "#### 📚 What you'll learn\n", @@ -33,7 +33,7 @@ }, { "cell_type": "markdown", - "id": "2f0a22a3", + "id": "9dc0f58c", "metadata": {}, "source": [ "### 📦 Import Data Designer\n", @@ -45,7 +45,7 @@ }, { "cell_type": "markdown", - "id": "9da71d7a", + "id": "0f0181d7", "metadata": {}, "source": [ "### ⚡ Colab Setup\n", @@ -56,7 +56,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a3c14c7b", + "id": "7a09d521", "metadata": {}, "outputs": [], "source": [ @@ -67,7 +67,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c5340c8f", + "id": "e7a60d1f", "metadata": {}, "outputs": [], "source": [ @@ -85,7 +85,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2906a904", + "id": "9348b4ca", "metadata": {}, "outputs": [], "source": [ @@ -108,7 +108,7 @@ }, { "cell_type": "markdown", - "id": "79670083", + "id": "875de027", "metadata": {}, "source": [ "### ⚙️ Initialize the Data Designer interface\n", @@ -121,7 +121,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ebb95a3e", + "id": "a926fe30", "metadata": {}, "outputs": [], "source": [ @@ -130,7 +130,7 @@ }, { "cell_type": "markdown", - "id": "e1f071c5", + "id": "ea2a1360", "metadata": {}, "source": [ "### 🏗️ Initialize the Data Designer Config Builder\n", @@ -145,7 +145,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4498ae34", + "id": "534b94e0", "metadata": {}, "outputs": [], "source": [ @@ -154,7 +154,7 @@ }, { "cell_type": "markdown", - "id": "7eb24bda", + "id": "4a90daaa", "metadata": {}, "source": [ "### 🌱 Seed Dataset Creation\n", @@ -171,7 +171,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8a667e92", + "id": "b4c5e452", "metadata": {}, "outputs": [], "source": [ @@ -186,7 +186,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b07af08d", + "id": "6fcf3c8b", "metadata": {}, "outputs": [], "source": [ @@ -231,7 +231,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f5e9a75b", + "id": "130a33c4", "metadata": {}, "outputs": [], "source": [ @@ -249,7 +249,7 @@ { "cell_type": "code", "execution_count": null, - "id": "194af0c7", + "id": "a017ba0a", "metadata": {}, "outputs": [], "source": [ @@ -259,7 +259,7 @@ { "cell_type": "code", "execution_count": null, - "id": "91f70616", + "id": "863de785", "metadata": {}, "outputs": [], "source": [ @@ -271,7 +271,7 @@ { "cell_type": "code", "execution_count": null, - "id": "688068c6", + "id": "1f4da220", "metadata": {}, "outputs": [], "source": [ @@ -293,7 +293,7 @@ }, { "cell_type": "markdown", - "id": "a7294c7d", + "id": "2af826c8", "metadata": {}, "source": [ "### 🔁 Iteration is key – preview the dataset!\n", @@ -310,7 +310,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b34374f9", + "id": "72479a9a", "metadata": {}, "outputs": [], "source": [ @@ -320,7 +320,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b787ead7", + "id": "e1938907", "metadata": {}, "outputs": [], "source": [ @@ -331,7 +331,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20bd1cc4", + "id": "17650be3", "metadata": {}, "outputs": [], "source": [ @@ -341,7 +341,7 @@ }, { "cell_type": "markdown", - "id": "0911e377", + "id": "3884bedb", "metadata": {}, "source": [ "### 📊 Analyze the generated data\n", @@ -354,7 +354,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fc09895a", + "id": "f720faa4", "metadata": {}, "outputs": [], "source": [ @@ -364,7 +364,7 @@ }, { "cell_type": "markdown", - "id": "49f093b5", + "id": "8d406105", "metadata": {}, "source": [ "### 🔎 Visual Inspection\n", @@ -375,7 +375,7 @@ { "cell_type": "code", "execution_count": null, - "id": "eaa91d52", + "id": "b4c6860c", "metadata": { "lines_to_next_cell": 2 }, @@ -399,7 +399,7 @@ }, { "cell_type": "markdown", - "id": "5c7db6e7", + "id": "b5be5452", "metadata": {}, "source": [ "### 🆙 Scale up!\n", @@ -412,7 +412,7 @@ { "cell_type": "code", "execution_count": null, - "id": "012913c5", + "id": "6f07a5ba", "metadata": {}, "outputs": [], "source": [ @@ -422,7 +422,7 @@ { "cell_type": "code", "execution_count": null, - "id": "027d6f26", + "id": "3a69f086", "metadata": {}, "outputs": [], "source": [ @@ -435,7 +435,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0e9803e1", + "id": "55dee3ba", "metadata": {}, "outputs": [], "source": [ @@ -447,7 +447,7 @@ }, { "cell_type": "markdown", - "id": "1ceba97d", + "id": "6f3dc05a", "metadata": {}, "source": [ "## ⏭️ Next Steps\n", diff --git a/docs/colab_notebooks/5-generating-images.ipynb b/docs/colab_notebooks/5-generating-images.ipynb index 83cab6d14..527822f2a 100644 --- a/docs/colab_notebooks/5-generating-images.ipynb +++ b/docs/colab_notebooks/5-generating-images.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "14e82e7c", + "id": "31e7b0a9", "metadata": {}, "source": [ "\"Open" @@ -10,7 +10,7 @@ }, { "cell_type": "markdown", - "id": "1cfe985b", + "id": "9b667aa6", "metadata": {}, "source": [ "# 🎨 Data Designer Tutorial: Generating Images\n", @@ -32,7 +32,7 @@ }, { "cell_type": "markdown", - "id": "c05d06db", + "id": "14778c16", "metadata": {}, "source": [ "### 📦 Import Data Designer\n", @@ -43,7 +43,7 @@ }, { "cell_type": "markdown", - "id": "5c9fdc69", + "id": "1d92676e", "metadata": {}, "source": [ "### ⚡ Colab Setup\n", @@ -54,7 +54,7 @@ { "cell_type": "code", "execution_count": null, - "id": "52f26a68", + "id": "ecfa7632", "metadata": {}, "outputs": [], "source": [ @@ -65,7 +65,7 @@ { "cell_type": "code", "execution_count": null, - "id": "983725dd", + "id": "ecee8556", "metadata": {}, "outputs": [], "source": [ @@ -83,7 +83,7 @@ { "cell_type": "code", "execution_count": null, - "id": "822bbf78", + "id": "e05085de", "metadata": {}, "outputs": [], "source": [ @@ -96,7 +96,7 @@ }, { "cell_type": "markdown", - "id": "3290a4a3", + "id": "1c6fd49c", "metadata": {}, "source": [ "### ⚙️ Initialize the Data Designer interface\n", @@ -107,7 +107,7 @@ { "cell_type": "code", "execution_count": null, - "id": "49ca144a", + "id": "6dd9e9a6", "metadata": {}, "outputs": [], "source": [ @@ -116,7 +116,7 @@ }, { "cell_type": "markdown", - "id": "fa690817", + "id": "63d457cc", "metadata": {}, "source": [ "### 🎛️ Define an image-generation model\n", @@ -128,7 +128,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e74fad29", + "id": "c2ab18c7", "metadata": {}, "outputs": [], "source": [ @@ -150,7 +150,7 @@ }, { "cell_type": "markdown", - "id": "695c04cb", + "id": "0f28d130", "metadata": {}, "source": [ "### 🏗️ Build the config: samplers + image column\n", @@ -161,7 +161,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0fdb0d8d", + "id": "acabaeb2", "metadata": {}, "outputs": [], "source": [ @@ -334,7 +334,7 @@ }, { "cell_type": "markdown", - "id": "3f835e6c", + "id": "a6c0c46f", "metadata": {}, "source": [ "### 🔁 Preview: images as base64\n", @@ -345,7 +345,7 @@ { "cell_type": "code", "execution_count": null, - "id": "eb54e0b9", + "id": "16df3f39", "metadata": {}, "outputs": [], "source": [ @@ -355,7 +355,7 @@ { "cell_type": "code", "execution_count": null, - "id": "61324ff2", + "id": "64e4e6b5", "metadata": {}, "outputs": [], "source": [ @@ -366,7 +366,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6a8d67d3", + "id": "cf41fad7", "metadata": {}, "outputs": [], "source": [ @@ -375,7 +375,7 @@ }, { "cell_type": "markdown", - "id": "7bd9070e", + "id": "0072df0b", "metadata": {}, "source": [ "### 🆙 Create: images saved to disk\n", @@ -386,7 +386,7 @@ { "cell_type": "code", "execution_count": null, - "id": "beabd69a", + "id": "4b641090", "metadata": {}, "outputs": [], "source": [ @@ -396,7 +396,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fa307c4c", + "id": "c2174264", "metadata": {}, "outputs": [], "source": [ @@ -407,7 +407,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a71c32d4", + "id": "549eaaa9", "metadata": {}, "outputs": [], "source": [ @@ -423,7 +423,7 @@ }, { "cell_type": "markdown", - "id": "7e1a1ea2", + "id": "309591ab", "metadata": {}, "source": [ "## ⏭️ Next steps\n", diff --git a/docs/colab_notebooks/6-editing-images-with-image-context.ipynb b/docs/colab_notebooks/6-editing-images-with-image-context.ipynb index 53dae31d1..121356c47 100644 --- a/docs/colab_notebooks/6-editing-images-with-image-context.ipynb +++ b/docs/colab_notebooks/6-editing-images-with-image-context.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "0eda27d4", + "id": "7d422a26", "metadata": {}, "source": [ "\"Open" @@ -10,7 +10,7 @@ }, { "cell_type": "markdown", - "id": "8a26005c", + "id": "9144792a", "metadata": {}, "source": [ "# 🎨 Data Designer Tutorial: Image-to-Image Editing\n", @@ -32,7 +32,7 @@ }, { "cell_type": "markdown", - "id": "3286587c", + "id": "fe623d18", "metadata": {}, "source": [ "### 📦 Import Data Designer\n", @@ -43,7 +43,7 @@ }, { "cell_type": "markdown", - "id": "05c1c87a", + "id": "f69a845c", "metadata": {}, "source": [ "### ⚡ Colab Setup\n", @@ -54,7 +54,7 @@ { "cell_type": "code", "execution_count": null, - "id": "636db7fc", + "id": "e191ac4c", "metadata": {}, "outputs": [], "source": [ @@ -65,7 +65,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a302e9fb", + "id": "dcff0dfa", "metadata": {}, "outputs": [], "source": [ @@ -83,7 +83,7 @@ { "cell_type": "code", "execution_count": null, - "id": "04ee6c29", + "id": "b9d5809a", "metadata": {}, "outputs": [], "source": [ @@ -99,7 +99,7 @@ }, { "cell_type": "markdown", - "id": "3f6e75cf", + "id": "4c0a7fbc", "metadata": {}, "source": [ "### ⚙️ Initialize the Data Designer interface\n", @@ -110,7 +110,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7979c23a", + "id": "808f7cde", "metadata": {}, "outputs": [], "source": [ @@ -119,7 +119,7 @@ }, { "cell_type": "markdown", - "id": "1723247a", + "id": "36c46907", "metadata": {}, "source": [ "### 🎛️ Define an image model\n", @@ -135,7 +135,7 @@ { "cell_type": "code", "execution_count": null, - "id": "234ddaeb", + "id": "a2ed9c86", "metadata": {}, "outputs": [], "source": [ @@ -157,7 +157,7 @@ }, { "cell_type": "markdown", - "id": "dc33263b", + "id": "648edf4c", "metadata": {}, "source": [ "### 🏗️ Build the configuration\n", @@ -172,7 +172,7 @@ { "cell_type": "code", "execution_count": null, - "id": "12d4691f", + "id": "da98219e", "metadata": {}, "outputs": [], "source": [ @@ -270,7 +270,7 @@ }, { "cell_type": "markdown", - "id": "eb2888aa", + "id": "3931d9f1", "metadata": {}, "source": [ "### 🔁 Preview: quick iteration\n", @@ -281,7 +281,7 @@ { "cell_type": "code", "execution_count": null, - "id": "293960ae", + "id": "63c63b91", "metadata": {}, "outputs": [], "source": [ @@ -291,7 +291,7 @@ { "cell_type": "code", "execution_count": null, - "id": "00d7cb2e", + "id": "38c99395", "metadata": {}, "outputs": [], "source": [ @@ -302,7 +302,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5a4c03e4", + "id": "5e5018e7", "metadata": {}, "outputs": [], "source": [ @@ -311,7 +311,7 @@ }, { "cell_type": "markdown", - "id": "0752cee9", + "id": "79322482", "metadata": { "lines_to_next_cell": 2 }, @@ -324,7 +324,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ae50ac45", + "id": "55a30a5c", "metadata": {}, "outputs": [], "source": [ @@ -355,7 +355,7 @@ { "cell_type": "code", "execution_count": null, - "id": "689c5222", + "id": "41e5d571", "metadata": {}, "outputs": [], "source": [ @@ -365,7 +365,7 @@ }, { "cell_type": "markdown", - "id": "b81eefbf", + "id": "f50c08ce", "metadata": {}, "source": [ "### 🆙 Create at scale\n", @@ -376,7 +376,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ea8d5ed1", + "id": "1c14d457", "metadata": {}, "outputs": [], "source": [ @@ -386,7 +386,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c0f45b1a", + "id": "8af80793", "metadata": {}, "outputs": [], "source": [ @@ -397,7 +397,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6e6f8001", + "id": "5aa0fe73", "metadata": {}, "outputs": [], "source": [ @@ -407,7 +407,7 @@ }, { "cell_type": "markdown", - "id": "4dc83d42", + "id": "03c37c62", "metadata": {}, "source": [ "## ⏭️ Next steps\n", diff --git a/docs/notebook_source/1-the-basics.py b/docs/notebook_source/1-the-basics.py index 8735d5821..f44f59d3b 100644 --- a/docs/notebook_source/1-the-basics.py +++ b/docs/notebook_source/1-the-basics.py @@ -324,7 +324,7 @@ # # Now that you've seen the basics of Data Designer, check out the following notebooks to learn more about: # -# - [Structured outputs and jinja expressions](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/2-structured-outputs-and-jinja-expressions/) +# - [Structured outputs, jinja expressions, and conditional generation](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/2-structured-outputs-and-jinja-expressions/) # # - [Seeding synthetic data generation with an external dataset](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/3-seeding-with-a-dataset/) # diff --git a/docs/notebook_source/2-structured-outputs-and-jinja-expressions.py b/docs/notebook_source/2-structured-outputs-and-jinja-expressions.py index df581612a..f389709f6 100644 --- a/docs/notebook_source/2-structured-outputs-and-jinja-expressions.py +++ b/docs/notebook_source/2-structured-outputs-and-jinja-expressions.py @@ -13,11 +13,11 @@ # --- # %% [markdown] -# # 🎨 Data Designer Tutorial: Structured Outputs and Jinja Expressions +# # 🎨 Data Designer Tutorial: Structured Outputs, Jinja Expressions, and Conditional Generation # # #### 📚 What you'll learn # -# In this notebook, we will continue our exploration of Data Designer, demonstrating more advanced data generation using structured outputs and Jinja expressions. +# In this notebook, we will continue our exploration of Data Designer, demonstrating more advanced data generation using structured outputs, Jinja expressions, and conditional generation with `skip.when`. # # If this is your first time using Data Designer, we recommend starting with the [first notebook](https://nvidia-nemo.github.io/DataDesigner/latest/notebooks/1-the-basics/) in this tutorial series. # @@ -305,6 +305,92 @@ class ProductReview(BaseModel): data_designer.validate(config_builder) +# %% [markdown] +# ## 🚦 Conditional generation with `skip.when` +# +# So far, every column is generated for every row. But sometimes an expensive LLM column only makes sense +# for a subset of rows — for example, a detailed complaint analysis is only useful when the review is negative. +# +# Data Designer lets you **skip** column generation on a per-row basis using `SkipConfig`. +# Skipped rows receive `None` by default, but you can provide a sentinel value with +# `skip=dd.SkipConfig(when="...", value="N/A")` to write a specific value instead. +# +# There are three patterns to know: +# +# | Pattern | How | Effect | +# |---|---|---| +# | **Expression gate** | `skip=dd.SkipConfig(when="...")` | Skip this column when the Jinja2 expression is truthy | +# | **Skip propagation** (default) | Downstream column depends on a skipped column | Automatically skipped too (`propagate_skip=True` by default) | +# | **Propagation opt-out** | `propagate_skip=False` on the downstream column | Always generates, even if an upstream was skipped | +# + +# %% [markdown] +# **Pattern 1 — Expression gate.** Only generate a detailed complaint analysis when the customer gave a low rating (1 or 2 stars). +# Rows where the rating is 3 or higher will get `None` for this column. +# + +# %% +config_builder.add_column( + dd.LLMTextColumnConfig( + name="complaint_analysis", + model_alias=MODEL_ALIAS, + prompt=( + "A customer reviewed '{{ product.name }}' ({{ product_category }} / {{ product_subcategory }}).\n\n" + "Review: {{ customer_review.review }}\n" + "Rating: {{ customer_review.rating }}/5\n" + "Mood: {{ customer_review.customer_mood }}\n\n" + "Write a short root-cause analysis of why this customer is unhappy " + "and suggest one concrete improvement the product team could make." + ), + skip=dd.SkipConfig(when="{{ customer_review.rating > 2 }}"), + ) +) + +# %% [markdown] +# **Pattern 2 — Skip propagation.** `action_items` depends on `complaint_analysis`. +# When `complaint_analysis` is skipped, `action_items` auto-skips too because +# `propagate_skip` defaults to `True`. +# + +# %% +config_builder.add_column( + dd.LLMTextColumnConfig( + name="action_items", + model_alias=MODEL_ALIAS, + prompt=( + "Based on this complaint analysis:\n" + "{{ complaint_analysis }}\n\n" + "List 2-3 concrete action items for the product team." + ), + ) +) + +# %% [markdown] +# **Pattern 3 — Propagation opt-out.** `review_summary` also depends on `complaint_analysis`, +# but sets `propagate_skip=False` so it always generates. The prompt uses a Jinja conditional +# to handle the case where `complaint_analysis` is `None`. +# + +# %% +config_builder.add_column( + dd.LLMTextColumnConfig( + name="review_summary", + model_alias=MODEL_ALIAS, + propagate_skip=False, + prompt=( + "Summarize this product review in one sentence:\n" + "Product: {{ product.name }}\n" + "Rating: {{ customer_review.rating }}/5\n" + "Review: {{ customer_review.review }}\n" + "{% if complaint_analysis %}" + "Complaint analysis: {{ complaint_analysis }}\n" + "{% endif %}" + ), + ) +) + +data_designer.validate(config_builder) + # %% [markdown] # ### 🔁 Iteration is key – preview the dataset! # @@ -322,10 +408,14 @@ class ProductReview(BaseModel): # %% # Run this cell multiple times to cycle through the 2 preview records. +# Look for rows where complaint_analysis and action_items are None (skipped) +# vs rows where they were generated (low-rated reviews). preview.display_sample_record() # %% # The preview dataset is available as a pandas DataFrame. +# Notice that complaint_analysis, action_items, and review_summary columns +# reflect the skip behavior: None for skipped rows, generated text otherwise. preview.dataset # %% [markdown] diff --git a/docs/notebook_source/_README.md b/docs/notebook_source/_README.md index 879e198f2..ff7dd541f 100644 --- a/docs/notebook_source/_README.md +++ b/docs/notebook_source/_README.md @@ -69,7 +69,7 @@ Learn the fundamentals of Data Designer by generating a simple product review da **Start here if you're new to Data Designer!** -### [2. Structured Outputs and Jinja Expressions](2-structured-outputs-and-jinja-expressions.ipynb) +### [2. Structured Outputs, Jinja Expressions, and Conditional Generation](2-structured-outputs-and-jinja-expressions.ipynb) Explore more advanced data generation capabilities: @@ -78,6 +78,7 @@ Explore more advanced data generation capabilities: - Combining samplers with structured data - Building complex data dependencies - Working with nested data structures +- Conditional generation with `skip.when` ### [3. Seeding with an External Dataset](3-seeding-with-a-dataset.ipynb) diff --git a/mkdocs.yml b/mkdocs.yml index fc4eaf383..bffbcc009 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,7 +34,7 @@ nav: - Tutorials: - Overview: notebooks/README.md - The Basics: notebooks/1-the-basics.ipynb - - Structured Outputs and Jinja Expressions: notebooks/2-structured-outputs-and-jinja-expressions.ipynb + - Structured Outputs, Jinja Expressions, and Conditional Generation: notebooks/2-structured-outputs-and-jinja-expressions.ipynb - Seeding with an External Dataset: notebooks/3-seeding-with-a-dataset.ipynb - Providing Images as Context: notebooks/4-providing-images-as-context.ipynb - Generating Images: notebooks/5-generating-images.ipynb diff --git a/packages/data-designer-config/src/data_designer/config/__init__.py b/packages/data-designer-config/src/data_designer/config/__init__.py index 72269519e..ed3336259 100644 --- a/packages/data-designer-config/src/data_designer/config/__init__.py +++ b/packages/data-designer-config/src/data_designer/config/__init__.py @@ -12,6 +12,7 @@ from data_designer.config.analysis.column_profilers import ( # noqa: F401 JudgeScoreProfilerConfig, ) + from data_designer.config.base import SkipConfig # noqa: F401 from data_designer.config.column_configs import ( # noqa: F401 CustomColumnConfig, EmbeddingColumnConfig, @@ -125,6 +126,8 @@ _LAZY_IMPORTS: dict[str, tuple[str, str]] = { # analysis.column_profilers "JudgeScoreProfilerConfig": (f"{_MOD_BASE}.analysis.column_profilers", "JudgeScoreProfilerConfig"), + # base + "SkipConfig": (f"{_MOD_BASE}.base", "SkipConfig"), # column_configs "CustomColumnConfig": (_MOD_COLUMN_CONFIGS, "CustomColumnConfig"), "EmbeddingColumnConfig": (_MOD_COLUMN_CONFIGS, "EmbeddingColumnConfig"), diff --git a/packages/data-designer-config/src/data_designer/config/base.py b/packages/data-designer-config/src/data_designer/config/base.py index 429d97a8f..31f0df571 100644 --- a/packages/data-designer-config/src/data_designer/config/base.py +++ b/packages/data-designer-config/src/data_designer/config/base.py @@ -6,8 +6,17 @@ from __future__ import annotations from abc import ABC, abstractmethod +from functools import cached_property -from pydantic import BaseModel, ConfigDict, Field +from jinja2 import meta as jinja2_meta +from jinja2.exceptions import TemplateSyntaxError +from jinja2.sandbox import ImmutableSandboxedEnvironment +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing_extensions import Self + +# Shared env for Jinja2 AST parsing (syntax checks + variable extraction). +# Cannot reuse misc.py helpers because base.py must not import data_designer.*. +_VALIDATION_ENV = ImmutableSandboxedEnvironment() class ConfigBase(BaseModel): @@ -20,6 +29,57 @@ class ConfigBase(BaseModel): ) +class SkipConfig(ConfigBase): + """Expression gate for conditional column generation. + + Attach to a ``SingleColumnConfig`` via ``skip=SkipConfig(...)`` to gate + generation on a Jinja2 expression. Controls *when* to skip; propagation + of upstream skips is controlled separately by ``propagate_skip`` on + ``SingleColumnConfig``. + + Attributes: + when: Jinja2 expression (including ``{{ }}`` delimiters); when truthy, + skip generation for this row. + value: Value to write for skipped cells. Defaults to ``None`` + (becomes ``NaN``/``pd.NA`` in the DataFrame). + """ + + when: str = Field( + description="Jinja2 expression (including {{ }} delimiters); when truthy, skip generation for this row.", + ) + value: bool | int | float | str | None = Field( + default=None, + description="Value to write for skipped cells. Defaults to None (becomes NaN/pd.NA in DataFrame).", + ) + + @field_validator("when") + @classmethod + def _validate_when_syntax(cls, v: str) -> str: + try: + ast = _VALIDATION_ENV.parse(v) + except TemplateSyntaxError as exc: + raise ValueError(str(exc)) from exc + if not jinja2_meta.find_undeclared_variables(ast): + raise ValueError( + f"skip.when expression {v!r} does not reference any columns. " + "Expressions must use Jinja2 delimiters, e.g. " + "'{{ in_stock == 0 }}' not 'in_stock == 0'." + ) + return v + + # cached_property writes to instance.__dict__; this works because ConfigBase + # is not frozen. If ConfigBase ever gains frozen=True, switch to model_post_init. + @cached_property + def columns(self) -> list[str]: + """Column names referenced in the ``when`` expression. + + Parsed once from the Jinja2 AST and cached. Used by the DAG builder + to add dependency edges and by the execution graph to store metadata. + """ + ast = _VALIDATION_ENV.parse(self.when) + return sorted(jinja2_meta.find_undeclared_variables(ast)) + + class SingleColumnConfig(ConfigBase, ABC): """Abstract base class for all single-column configuration types. @@ -32,12 +92,47 @@ class SingleColumnConfig(ConfigBase, ABC): Useful for intermediate columns that are dependencies for other columns. column_type: Discriminator field that identifies the specific column type. Subclasses must override this field to specify the column type with a `Literal` value. + skip: Optional expression gate for conditional generation. + propagate_skip: If True (default), this column auto-skips when any of its + required_columns was skipped. Independent of ``skip``. """ name: str drop: bool = False allow_resize: bool = False column_type: str + skip: SkipConfig | None = None + propagate_skip: bool = Field( + default=True, + description="If True (default), this column auto-skips when any " + "of its required_columns was skipped. Independent of skip — " + "a column with no SkipConfig still propagates upstream skips. " + "Set to False for null-tolerant columns.", + ) + + @model_validator(mode="after") + def _validate_skip_scope(self) -> Self: + if self.skip is not None: + if self.column_type in ("sampler", "seed-dataset"): + raise ValueError( + f"skip is not supported on {self.column_type} columns. " + "Sampler/seed columns are collapsed into shared multi-column generators " + "and cannot be skipped individually." + ) + if self.allow_resize: + raise ValueError( + "skip and allow_resize cannot be used together. " + "allow_resize changes buffer size during generation (1:N / N:1), which " + "breaks index-based skip tracking and merge-back." + ) + self_refs = {self.name} | set(self.side_effect_columns) + if not self_refs.isdisjoint(self.skip.columns): + offending = self_refs & set(self.skip.columns) + raise ValueError( + f"skip.when expression for column '{self.name}' references itself " + f"(via {offending!r}). A column cannot gate its own generation." + ) + return self @staticmethod def get_column_emoji() -> str: diff --git a/packages/data-designer-config/src/data_designer/config/column_configs.py b/packages/data-designer-config/src/data_designer/config/column_configs.py index 59bd9e39a..88dffe9a4 100644 --- a/packages/data-designer-config/src/data_designer/config/column_configs.py +++ b/packages/data-designer-config/src/data_designer/config/column_configs.py @@ -415,8 +415,15 @@ def required_columns(self) -> list[str]: def side_effect_columns(self) -> list[str]: return [] + _DTYPE_COERCERS: dict[str, type] = { + "int": int, + "float": float, + "str": str, + "bool": bool, + } + @model_validator(mode="after") - def assert_expression_valid_jinja(self) -> Self: + def _assert_expression_valid_jinja(self) -> Self: """Validate that the expression is a valid, non-empty Jinja2 template. Returns: @@ -434,6 +441,22 @@ def assert_expression_valid_jinja(self) -> Self: assert_valid_jinja2_template(self.expr) return self + @model_validator(mode="after") + def _coerce_skip_value_to_dtype(self) -> Self: + """Coerce ``skip.value`` to match ``dtype`` so skipped and computed rows share a type.""" + if self.skip is None or self.skip.value is None: + return self + target_type = self._DTYPE_COERCERS.get(self.dtype) + if target_type is not None and not isinstance(self.skip.value, target_type): + try: + self.skip.value = target_type(self.skip.value) + except (ValueError, TypeError) as exc: + raise ValueError( + f"Expression column '{self.name}' has dtype='{self.dtype}' but " + f"skip.value={self.skip.value!r} cannot be converted to {self.dtype}: {exc}" + ) from exc + return self + class ValidationColumnConfig(SingleColumnConfig): """Configuration for validation columns that validate existing columns. diff --git a/packages/data-designer-config/tests/config/test_columns.py b/packages/data-designer-config/tests/config/test_columns.py index 987a158f6..0c166578b 100644 --- a/packages/data-designer-config/tests/config/test_columns.py +++ b/packages/data-designer-config/tests/config/test_columns.py @@ -6,6 +6,7 @@ import pytest from pydantic import ValidationError +from data_designer.config.base import SkipConfig from data_designer.config.column_configs import ( EmbeddingColumnConfig, ExpressionColumnConfig, @@ -600,3 +601,47 @@ def test_allow_resize_inherited_by_subclasses() -> None: """Subclasses inherit allow_resize from SingleColumnConfig.""" assert StubColumnConfig(name="test").allow_resize is False assert StubColumnConfig(name="test", allow_resize=True).allow_resize is True + + +@pytest.mark.parametrize( + ("dtype", "raw_value", "expected_value", "expected_type"), + [ + pytest.param("float", 0, 0.0, float, id="int-to-float"), + pytest.param("int", 0.0, 0, int, id="float-to-int"), + pytest.param("str", 0.0, "0.0", str, id="float-to-str"), + pytest.param("str", 42, "42", str, id="int-to-str"), + pytest.param("bool", 1, True, bool, id="int-to-bool"), + pytest.param("float", "3.14", 3.14, float, id="str-to-float"), + ], +) +def test_expression_column_skip_value_coerced_to_dtype( + dtype: str, raw_value: object, expected_value: object, expected_type: type +) -> None: + config = ExpressionColumnConfig( + name="col", + expr="{{ price }}", + dtype=dtype, + skip=SkipConfig(when="{{ flag }}", value=raw_value), + ) + assert config.skip.value == expected_value + assert type(config.skip.value) is expected_type + + +def test_expression_column_skip_value_none_left_alone() -> None: + config = ExpressionColumnConfig( + name="col", + expr="{{ price }}", + dtype="float", + skip=SkipConfig(when="{{ flag }}"), + ) + assert config.skip.value is None + + +def test_expression_column_skip_value_incompatible_raises() -> None: + with pytest.raises(ValidationError, match="cannot be converted"): + ExpressionColumnConfig( + name="col", + expr="{{ price }}", + dtype="int", + skip=SkipConfig(when="{{ flag }}", value="not-a-number"), + ) diff --git a/packages/data-designer-config/tests/config/test_skip_config.py b/packages/data-designer-config/tests/config/test_skip_config.py new file mode 100644 index 000000000..6b94fd21f --- /dev/null +++ b/packages/data-designer-config/tests/config/test_skip_config.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from data_designer.config.base import SkipConfig +from data_designer.config.column_configs import ( + LLMTextColumnConfig, + SamplerColumnConfig, + SeedDatasetColumnConfig, +) +from data_designer.config.sampler_params import SamplerType, UUIDSamplerParams + +_BASE_LLM = dict(name="test", prompt="test {{ x }}", model_alias="default") + + +@pytest.mark.parametrize( + "value", + [True, False, 0, 42, 1.5, "skipped", None], +) +def test_skip_config_value_types(value: bool | int | float | str | None) -> None: + cfg = SkipConfig(when="{{ x == 0 }}", value=value) + assert cfg.value == value + + +@pytest.mark.parametrize( + ("when", "match"), + [ + pytest.param("{{ 1 + }}", "unexpected", id="syntax-error"), + pytest.param("in_stock == 0", "does not reference any columns", id="no-delimiters"), + ], +) +def test_skip_config_when_rejects_invalid_expressions(when: str, match: str) -> None: + with pytest.raises(ValidationError, match=match): + SkipConfig(when=when) + + +@pytest.mark.parametrize( + ("when", "expected"), + [ + pytest.param("{{ in_stock == 0 }}", ["in_stock"], id="single"), + pytest.param("{{ a > 0 and b < 10 }}", ["a", "b"], id="multiple"), + ], +) +def test_skip_config_columns_extraction(when: str, expected: list[str]) -> None: + cfg = SkipConfig(when=when) + assert cfg.columns == expected + + +def test_skip_config_columns_cached() -> None: + cfg = SkipConfig(when="{{ x == 0 }}") + first = cfg.columns + assert cfg.columns is first + + +@pytest.mark.parametrize( + ("attr", "expected"), + [ + pytest.param("skip", None, id="skip-defaults-none"), + pytest.param("propagate_skip", True, id="propagate-skip-defaults-true"), + ], +) +def test_single_column_config_skip_defaults(attr: str, expected: object) -> None: + col = LLMTextColumnConfig(**_BASE_LLM) + assert getattr(col, attr) == expected + + +def test_skip_rejected_on_sampler_type() -> None: + with pytest.raises(ValidationError, match="skip is not supported on sampler columns"): + SamplerColumnConfig( + name="s", + sampler_type=SamplerType.UUID, + params=UUIDSamplerParams(prefix="p_", short_form=True), + skip=SkipConfig(when="{{ y == 1 }}"), + ) + + +def test_skip_rejected_on_seed_dataset_type() -> None: + with pytest.raises(ValidationError, match="skip is not supported on seed-dataset columns"): + SeedDatasetColumnConfig( + name="seed_col", + skip=SkipConfig(when="{{ y == 1 }}"), + ) + + +def test_skip_rejected_with_allow_resize() -> None: + with pytest.raises(ValidationError, match="skip and allow_resize cannot be used together"): + LLMTextColumnConfig( + **_BASE_LLM, + allow_resize=True, + skip=SkipConfig(when="{{ x == 0 }}"), + ) + + +def test_skip_self_reference_rejected() -> None: + with pytest.raises(ValidationError, match="references itself"): + LLMTextColumnConfig( + name="foo", + prompt="test {{ bar }}", + model_alias="default", + skip=SkipConfig(when="{{ foo == 0 }}"), + ) + + +def test_skip_side_effect_self_reference_rejected() -> None: + """Referencing a column's own side-effect (e.g. trace) in skip.when is a self-reference.""" + from data_designer.config.column_configs import TraceType + + with pytest.raises(ValidationError, match="references itself"): + LLMTextColumnConfig( + name="review", + prompt="test {{ bar }}", + model_alias="default", + with_trace=TraceType.ALL_MESSAGES, + skip=SkipConfig(when="{{ review__trace == 'x' }}"), + ) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py index 607b8f89f..067795363 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py @@ -15,12 +15,18 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.config.column_configs import GenerationStrategy from data_designer.engine.context import current_row_group +from data_designer.engine.dataset_builders.multi_column_configs import MultiColumnConfig from data_designer.engine.dataset_builders.utils.async_progress_reporter import ( DEFAULT_REPORT_INTERVAL, AsyncProgressReporter, ) from data_designer.engine.dataset_builders.utils.completion_tracker import CompletionTracker from data_designer.engine.dataset_builders.utils.progress_tracker import ProgressTracker +from data_designer.engine.dataset_builders.utils.skip_evaluator import should_skip_column_for_record +from data_designer.engine.dataset_builders.utils.skip_tracker import ( + apply_skip_to_record, + strip_skip_metadata_from_records, +) from data_designer.engine.dataset_builders.utils.sticky_progress_bar import StickyProgressBar from data_designer.engine.dataset_builders.utils.task_model import SliceRef, Task, TaskTrace from data_designer.engine.models.errors import ( @@ -704,10 +710,11 @@ async def _execute_task_inner_impl(self, task: Task) -> None: if self._trace and trace: trace.slot_acquired_at = time.perf_counter() + cell_skipped = False if task.task_type == "from_scratch": await self._run_from_scratch(task, generator) elif task.task_type == "cell": - await self._run_cell(task, generator) + _result, cell_skipped = await self._run_cell(task, generator) elif task.task_type == "batch": await self._run_batch(task, generator) else: @@ -723,7 +730,10 @@ async def _execute_task_inner_impl(self, task: Task) -> None: self._check_error_rate(success=True) if self._reporter: - self._reporter.record_success(task.column) + if cell_skipped: + self._reporter.record_skipped(task.column) + else: + self._reporter.record_success(task.column) if self._trace and trace: trace.status = "ok" @@ -788,21 +798,29 @@ async def _run_from_scratch(self, task: Task, generator: ColumnGenerator) -> Any return result_df - async def _run_cell(self, task: Task, generator: ColumnGenerator) -> Any: - """Execute a cell-by-cell task.""" + async def _run_cell(self, task: Task, generator: ColumnGenerator) -> tuple[Any, bool]: + """Execute a cell-by-cell task. Returns ``(result, skipped)``.""" if task.row_index is None: raise ValueError(f"Cell task requires a row_index, got None for column '{task.column}'") if self._tracker.is_dropped(task.row_group, task.row_index): - return None + return None, False - # Read row from buffer + # Evaluate skip against the live buffer record (no copy needed — + # there is no `await` between the read and the skip-metadata write). if self._buffer_manager is not None: - row_data = dict(self._buffer_manager.get_row(task.row_group, task.row_index)) + record = self._buffer_manager.get_row(task.row_group, task.row_index) else: - row_data = {} + record = {} - result = await generator.agenerate(row_data) + if self._should_skip_record(task.column, record): + self._apply_skip_to_record(task, record) + skip_config = self._graph.get_skip_config(task.column) + return skip_config.value if skip_config is not None else None, True + + # Copy for generation: agenerate crosses an await boundary, so the + # generator must not hold a mutable reference to the live record. + result = await generator.agenerate(dict(record)) # Write back to buffer (include side-effect columns) if self._buffer_manager is not None and not self._tracker.is_dropped(task.row_group, task.row_index): @@ -811,27 +829,73 @@ async def _run_cell(self, task: Task, generator: ColumnGenerator) -> Any: if col in result: self._buffer_manager.update_cell(task.row_group, task.row_index, col, result[col]) - return result + return result, False + + def _should_skip_record(self, column: str, record: dict) -> bool: + """Decide whether a cell should be skipped (propagation first, then expression gate).""" + skip_config = self._graph.get_skip_config(column) + return should_skip_column_for_record( + record, + propagate_skip=self._graph.should_propagate_skip(column), + required_columns=self._graph.get_required_columns(column), + skip_config_when=skip_config.when if skip_config is not None else None, + ) + + def _apply_skip_to_record(self, task: Task, record: dict) -> None: + """Write skip metadata directly into *record* (the live buffer row).""" + skip_config = self._graph.get_skip_config(task.column) + skip_value = skip_config.value if skip_config is not None else None + apply_skip_to_record( + record, + column_name=task.column, + cell_value=skip_value, + side_effect_columns=self._graph.get_side_effect_columns(task.column), + ) async def _run_batch(self, task: Task, generator: ColumnGenerator) -> Any: """Execute a full-column/batch task.""" + rg_size = self._get_rg_size(task.row_group) + if self._buffer_manager is not None: - batch_df = self._buffer_manager.get_dataframe(task.row_group) - # Snapshot dropped rows before the await so the row-count expectation - # is consistent with batch_df (concurrent tasks may drop rows during agenerate). - rg_size = self._get_rg_size(task.row_group) pre_dropped: set[int] = {ri for ri in range(rg_size) if self._buffer_manager.is_dropped(task.row_group, ri)} + active_rows_data: list[dict] = [] + + # Skip evaluation only applies to single-column configs. + # Multi-column configs (sampler/seed) are rejected by the SkipConfig + # model validator, so they never carry skip metadata. + pre_skipped: set[int] = set() + is_multi = isinstance(generator.config, MultiColumnConfig) + for ri in range(rg_size): + if ri in pre_dropped: + continue + + record = self._buffer_manager.get_row(task.row_group, ri) + if not is_multi and self._should_skip_record(task.column, record): + self._apply_skip_to_record(task, record) + pre_skipped.add(ri) + continue + + active_rows_data.append(record) + + batch_df = ( + lazy.pd.DataFrame(strip_skip_metadata_from_records(active_rows_data)) + if active_rows_data + else lazy.pd.DataFrame() + ) else: batch_df = lazy.pd.DataFrame() - rg_size = self._get_rg_size(task.row_group) pre_dropped = set() + pre_skipped = set() + + if len(batch_df) == 0: + return batch_df result_df = await generator.agenerate(batch_df) # Merge result columns back to buffer (include side-effect columns) if self._buffer_manager is not None: write_cols = self._gen_instance_to_columns_including_side_effects.get(id(generator), [task.column]) - active_rows = rg_size - len(pre_dropped) + active_rows = rg_size - len(pre_dropped) - len(pre_skipped) if len(result_df) != active_rows: raise ValueError( f"Batch generator for '{task.column}' returned {len(result_df)} rows " @@ -839,9 +903,8 @@ async def _run_batch(self, task: Task, generator: ColumnGenerator) -> Any: ) result_idx = 0 for ri in range(rg_size): - if ri in pre_dropped: + if ri in pre_dropped or ri in pre_skipped: continue - # Skip writing to rows dropped by concurrent tasks during the await if not self._buffer_manager.is_dropped(task.row_group, ri): for col in write_cols: if col in result_df.columns: diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py index f97054340..182f0438e 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py @@ -35,8 +35,17 @@ from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager +from data_designer.engine.dataset_builders.utils.execution_graph import ExecutionGraph from data_designer.engine.dataset_builders.utils.processor_runner import ProcessorRunner, ProcessorStage from data_designer.engine.dataset_builders.utils.progress_tracker import ProgressTracker +from data_designer.engine.dataset_builders.utils.skip_evaluator import should_skip_column_for_record +from data_designer.engine.dataset_builders.utils.skip_tracker import ( + SKIPPED_COLUMNS_RECORD_KEY, + apply_skip_to_record, + prepare_records_for_skip_metadata_round_trip, + restore_skip_metadata, + strip_skip_metadata_from_records, +) from data_designer.engine.dataset_builders.utils.sticky_progress_bar import StickyProgressBar from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler from data_designer.engine.processing.processors.base import Processor @@ -76,7 +85,6 @@ ensure_async_engine_loop, ) from data_designer.engine.dataset_builders.utils.completion_tracker import CompletionTracker - from data_designer.engine.dataset_builders.utils.execution_graph import ExecutionGraph from data_designer.engine.dataset_builders.utils.row_group_buffer import RowGroupBufferManager @@ -97,6 +105,7 @@ def __init__( self._cell_resize_mode = False self._task_traces: list[TaskTrace] = [] self._registry = registry or DataDesignerRegistry() + self._graph: ExecutionGraph | None = None self._data_designer_config = compile_data_designer_config(data_designer_config, resource_provider) self._column_configs = compile_dataset_builder_column_configs(self._data_designer_config) @@ -136,6 +145,10 @@ def single_column_configs(self) -> list[ColumnConfigT]: configs.append(config) return configs + @functools.cached_property + def single_column_config_by_name(self) -> dict[str, ColumnConfigT]: + return {config.name: config for config in self.single_column_configs} + @functools.cached_property def llm_generated_column_configs(self) -> list[ColumnConfigT]: return [config for config in self.single_column_configs if column_type_is_model_generated(config.column_type)] @@ -168,7 +181,7 @@ def build( mode = StorageMode.DISK if save_multimedia_to_disk else StorageMode.DATAFRAME self.artifact_storage.set_media_storage_mode(mode) - generators = self._initialize_generators() + generators, self._graph = self._initialize_generators_and_graph() start_time = time.perf_counter() buffer_size = self._resource_provider.run_config.buffer_size @@ -202,7 +215,7 @@ def build_preview(self, *, num_records: int) -> pd.DataFrame: if self._has_image_columns(): self.artifact_storage.set_media_storage_mode(StorageMode.DATAFRAME) - generators = self._initialize_generators() + generators, self._graph = self._initialize_generators_and_graph() start_time = time.perf_counter() if DATA_DESIGNER_ASYNC_ENGINE: @@ -409,13 +422,23 @@ def _has_image_columns(self) -> bool: """Check if config has any image generation columns.""" return any(col.column_type == DataDesignerColumnType.IMAGE for col in self.single_column_configs) - def _initialize_generators(self) -> list[ColumnGenerator]: - return [ + def _initialize_generators_and_graph(self) -> tuple[list[ColumnGenerator], ExecutionGraph]: + generators = [ self._registry.column_generators.get_for_config_type(type(config))( config=config, resource_provider=self._resource_provider ) for config in self._column_configs ] + strategies: dict[str, GenerationStrategy] = {} + for gen in generators: + strategy = gen.get_generation_strategy() + if isinstance(gen.config, MultiColumnConfig): + for sub in gen.config.columns: + strategies[sub.name] = strategy + else: + strategies[gen.config.name] = strategy + graph = ExecutionGraph.create(self._column_configs, strategies) + return generators, graph def _write_builder_config(self) -> None: self.artifact_storage.mkdir_if_needed(self.artifact_storage.base_dataset_path) @@ -500,12 +523,137 @@ def _log_resize_if_changed(self, column_name: str, original_count: int, new_coun emoji = "💥" if new_count > original_count else "✂️" logger.info(f"{emoji} Column '{column_name}' resized batch: {original_count} -> {new_count} records.") + def _require_graph(self) -> ExecutionGraph: + """Return the initialized execution graph for the current run.""" + graph = self._graph + if graph is None: + raise DatasetGenerationError("Execution graph accessed before generator initialization.") + return graph + + def _column_can_skip(self, column_name: str) -> bool: + """Fast check: can *column_name* ever be skipped (expression gate or propagation)? + + Returns ``False`` for ``allow_resize=True`` columns because 1:N generators + change the row count — the skip-aware merge path assumes a 1:1 mapping + between input and output rows and would raise on the row-count check. + """ + if self._graph is None: + return False + config = self.single_column_config_by_name.get(column_name) + if config is not None and config.allow_resize: + return False + if self._graph.get_skip_config(column_name) is not None: + return True + return self._graph.should_propagate_skip(column_name) and bool(self._graph.get_required_columns(column_name)) + + def _should_skip_cell(self, column_name: str, record: dict) -> bool: + """Decide whether a single cell should be skipped (propagation or expression gate).""" + skip_config = self._graph.get_skip_config(column_name) + return should_skip_column_for_record( + record, + propagate_skip=self._graph.should_propagate_skip(column_name), + required_columns=self._graph.get_required_columns(column_name), + skip_config_when=skip_config.when if skip_config is not None else None, + ) + + def _write_skip_to_record(self, column_name: str, record: dict) -> None: + """Write skip metadata and the skip value into *record* in-place.""" + skip_config = self._graph.get_skip_config(column_name) + skip_value = skip_config.value if skip_config is not None else None + apply_skip_to_record( + record, + column_name=column_name, + cell_value=skip_value, + side_effect_columns=self._graph.get_side_effect_columns(column_name), + ) + def _run_full_column_generator(self, generator: ColumnGenerator) -> None: + column_name = generator.config.name if not isinstance(generator.config, MultiColumnConfig) else None + + if column_name is not None and self._column_can_skip(column_name): + self._run_full_column_generator_with_skip(generator, column_name) + else: + self._run_full_column_generator_without_skip(generator) + + def _run_full_column_generator_without_skip(self, generator: ColumnGenerator) -> None: + """Run the generator on the full batch, preserving skip metadata across the replace.""" original_count = self.batch_manager.num_records_in_buffer - df = generator.generate(self.batch_manager.get_current_batch(as_dataframe=True)) - allow_resize = getattr(generator.config, "allow_resize", False) + allow_resize = generator.config.allow_resize if not isinstance(generator.config, MultiColumnConfig) else False + old_records = [record for _, record in self.batch_manager.iter_current_batch()] + input_records, restore_context = prepare_records_for_skip_metadata_round_trip(old_records) + + df = generator.generate(lazy.pd.DataFrame(input_records)) self._log_resize_if_changed(self._column_display_name(generator.config), original_count, len(df), allow_resize) - self.batch_manager.replace_buffer(df.to_dict(orient="records"), allow_resize=allow_resize) + new_records = df.to_dict(orient="records") + if restore_context is not None: + try: + restore_skip_metadata(new_records, context=restore_context, allow_resize=allow_resize) + except ValueError as exc: + raise DatasetGenerationError( + f"Unable to restore skip provenance after FULL_COLUMN generation for " + f"{self._column_display_name(generator.config)}: {exc}" + ) from exc + self.batch_manager.replace_buffer(new_records, allow_resize=allow_resize) + + def _run_full_column_generator_with_skip(self, generator: ColumnGenerator, column_name: str) -> None: + """Run a FULL_COLUMN generator with per-row skip evaluation and merge-back. + + Only reachable when ``_column_can_skip`` is True, which excludes + ``allow_resize=True`` columns, so resize handling is not needed here. + """ + active_records: list[dict] = [] + records_with_skip_status: list[tuple[bool, dict]] = [] + has_skipped = False + for _, record in self.batch_manager.iter_current_batch(): + skipped = self._should_skip_cell(column_name, record) + if skipped: + has_skipped = True + self._write_skip_to_record(column_name, record) + else: + active_records.append(record) + records_with_skip_status.append((skipped, record)) + + if not has_skipped: + # No rows were actually skipped — use the normal path to avoid the + # overhead of stripping metadata, building a separate active DataFrame, + # and merging results back. + self._run_full_column_generator_without_skip(generator) + return + + batch = self._merge_skipped_and_generated(generator, column_name, active_records, records_with_skip_status) + self.batch_manager.replace_buffer(batch, allow_resize=False) + + def _merge_skipped_and_generated( + self, + generator: ColumnGenerator, + column_name: str, + active_records: list[dict], + records_with_skip_status: list[tuple[bool, dict]], + ) -> list[dict]: + """Generate only for active (non-skipped) records and merge back with skipped ones.""" + if not active_records: + return [record for _, record in records_with_skip_status] + + active_df = lazy.pd.DataFrame(strip_skip_metadata_from_records(active_records)) + result_records = generator.generate(active_df).to_dict(orient="records") + if len(result_records) != len(active_records): + raise DatasetGenerationError( + f"Generator for '{column_name}' returned {len(result_records)} rows " + f"but {len(active_records)} active (non-skipped) records were expected." + ) + + result_iter = iter(result_records) + batch: list[dict] = [] + for skipped, record in records_with_skip_status: + if skipped: + batch.append(record) + continue + gen_result = next(result_iter) + prior_skipped = record.get(SKIPPED_COLUMNS_RECORD_KEY) + if prior_skipped is not None: + gen_result[SKIPPED_COLUMNS_RECORD_KEY] = prior_skipped + batch.append(gen_result) + return batch def _run_model_health_check_if_needed(self) -> None: model_aliases: set[str] = set() @@ -615,16 +763,23 @@ def _fan_out_with_async(self, generator: ColumnGeneratorWithModelRegistry, max_w if getattr(generator.config, "tool_alias", None): logger.info("🛠️ Tool calling enabled") bar = StickyProgressBar() if self._resource_provider.run_config.progress_bar else None + can_skip = self._column_can_skip(generator.config.name) with bar or contextlib.nullcontext(): progress_tracker, executor_kwargs = self._setup_fan_out(generator, max_workers, progress_bar=bar) executor = AsyncConcurrentExecutor(max_workers=max_workers, **executor_kwargs) - work_items = [ - ( - generator.agenerate(record), - {"index": i, "column_name": generator.config.name}, + work_items: list[tuple[Any, dict[str, Any]]] = [] + for i, record in self.batch_manager.iter_current_batch(): + if can_skip and self._should_skip_cell(generator.config.name, record): + self._write_skip_to_record(generator.config.name, record) + self.batch_manager.update_record(i, record) + progress_tracker.record_skipped() + continue + work_items.append( + ( + generator.agenerate(record), + {"index": i, "column_name": generator.config.name}, + ) ) - for i, record in self.batch_manager.iter_current_batch() - ] executor.run(work_items) self._finalize_fan_out(progress_tracker) @@ -632,10 +787,16 @@ def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max if getattr(generator.config, "tool_alias", None): logger.info("🛠️ Tool calling enabled") bar = StickyProgressBar() if self._resource_provider.run_config.progress_bar else None + can_skip = self._column_can_skip(generator.config.name) with bar or contextlib.nullcontext(): progress_tracker, executor_kwargs = self._setup_fan_out(generator, max_workers, progress_bar=bar) with ConcurrentThreadExecutor(max_workers=max_workers, **executor_kwargs) as executor: for i, record in self.batch_manager.iter_current_batch(): + if can_skip and self._should_skip_cell(generator.config.name, record): + self._write_skip_to_record(generator.config.name, record) + self.batch_manager.update_record(i, record) + progress_tracker.record_skipped() + continue executor.submit( lambda record: generator.generate(record), record, diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py index 60e5583e8..4b3e03670 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py @@ -4,6 +4,7 @@ from __future__ import annotations import logging +from itertools import chain import data_designer.lazy_heavy_imports as lazy from data_designer.config.column_types import ColumnConfigT @@ -28,6 +29,7 @@ def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> li return non_dag_column_config_list side_effect_dict = {n: list(c.side_effect_columns) for n, c in dag_column_config_dict.items()} + all_side_effects = set(chain.from_iterable(side_effect_dict.values())) side_effect_to_producer: dict[str, str] = {} for producer, cols in side_effect_dict.items(): @@ -44,19 +46,13 @@ def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> li logger.info("⛓️ Sorting column configs into a Directed Acyclic Graph") for name, col in dag_column_config_dict.items(): dag.add_node(name) - for req_col_name in col.required_columns: - if req_col_name in list(dag_column_config_dict.keys()): - logger.debug(f"{LOG_INDENT}🔗 `{name}` depends on `{req_col_name}`") - dag.add_edge(req_col_name, name) - - # If the required column is a side effect of another column, - # add an edge from the parent column to the current column. - elif req_col_name in sum(side_effect_dict.values(), []): - for parent, cols in side_effect_dict.items(): - if req_col_name in cols: - logger.debug(f"{LOG_INDENT}🔗 `{name}` depends on `{parent}` via `{req_col_name}`") - dag.add_edge(parent, name) - break + _add_dependency_edges( + dag, name, list(col.required_columns), dag_column_config_dict, side_effect_dict, all_side_effects, "" + ) + if col.skip is not None: + _add_dependency_edges( + dag, name, col.skip.columns, dag_column_config_dict, side_effect_dict, all_side_effects, "skip.when" + ) if not lazy.nx.is_directed_acyclic_graph(dag): raise DAGCircularDependencyError( @@ -69,3 +65,25 @@ def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> li sorted_columns.extend([dag_column_config_dict[n] for n in list(lazy.nx.topological_sort(dag))]) return sorted_columns + + +def _add_dependency_edges( + dag: lazy.nx.DiGraph, + name: str, + dep_names: list[str], + dag_column_config_dict: dict[str, ColumnConfigT], + side_effect_dict: dict[str, list[str]], + all_side_effects: set[str], + label: str, +) -> None: + """Add DAG edges from *dep_names* to *name*, resolving through side-effect parents.""" + for dep in dep_names: + if dep in dag_column_config_dict: + logger.debug(f"{LOG_INDENT}🔗 `{name}` {label} depends on `{dep}`") + dag.add_edge(dep, name) + elif dep in all_side_effects: + for parent, cols in side_effect_dict.items(): + if dep in cols: + logger.debug(f"{LOG_INDENT}🔗 `{name}` {label} depends on `{parent}` via `{dep}`") + dag.add_edge(parent, name) + break diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dataset_batch_manager.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dataset_batch_manager.py index 09e6ce9c8..757853300 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dataset_batch_manager.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dataset_batch_manager.py @@ -10,6 +10,7 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.engine.dataset_builders.utils.errors import DatasetBatchManagementError +from data_designer.engine.dataset_builders.utils.skip_tracker import strip_skip_metadata_from_records from data_designer.engine.storage.artifact_storage import ArtifactStorage, BatchStage if TYPE_CHECKING: @@ -133,7 +134,7 @@ def get_current_batch_number(self) -> int: def get_current_batch(self, *, as_dataframe: bool = False) -> pd.DataFrame | list[dict]: if as_dataframe: - return lazy.pd.DataFrame(self._buffer) + return lazy.pd.DataFrame(strip_skip_metadata_from_records(self._buffer)) return self._buffer def iter_current_batch(self) -> Iterator[tuple[int, dict]]: @@ -182,7 +183,7 @@ def write(self) -> Path | None: try: file_path = self.artifact_storage.write_batch_to_parquet_file( batch_number=self._current_batch_number, - dataframe=lazy.pd.DataFrame(self._buffer), + dataframe=lazy.pd.DataFrame(strip_skip_metadata_from_records(self._buffer)), batch_stage=BatchStage.PARTIAL_RESULT, ) return file_path diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py index cbf8cf104..5cd41dd38 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py @@ -6,6 +6,7 @@ import logging import math from collections import deque +from typing import TYPE_CHECKING from data_designer.config.column_configs import GenerationStrategy from data_designer.engine.dataset_builders.multi_column_configs import ( @@ -17,6 +18,9 @@ logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from data_designer.config.base import SkipConfig + class ExecutionGraph: """Column-level static execution graph built from column configs. @@ -34,6 +38,10 @@ def __init__(self) -> None: self._columns: list[str] = [] self._topological_order_cache: list[str] | None = None self._upstream_by_strategy_cache: dict[str, tuple[list[str], list[str]]] = {} + self._required_columns: dict[str, list[str]] = {} + self._skip_configs: dict[str, SkipConfig] = {} + self._propagate_skip: dict[str, bool] = {} + self._producer_to_side_effect_map: dict[str, list[str]] = {} @property def columns(self) -> list[str]: @@ -55,7 +63,7 @@ def create( """ graph = cls() - # First pass: register all columns, strategies, and side-effect mappings + # First pass: register all columns, strategies, side-effect mappings, and skip metadata for config in column_configs: if isinstance(config, MultiColumnConfig): sub_configs = config.columns @@ -69,9 +77,13 @@ def create( for se_col in sub.side_effect_columns: graph.set_side_effect(se_col, name) + graph.set_propagate_skip(name, sub.propagate_skip) + if sub.skip is not None: + graph.set_skip_config(name, sub.skip) + known_columns = set(graph.columns) - # Second pass: build edges + # Second pass: build edges (required_columns + skip.columns) for config in column_configs: if isinstance(config, MultiColumnConfig): sub_configs = config.columns @@ -80,6 +92,7 @@ def create( for sub in sub_configs: name = sub.name + resolved_required: list[str] = [] for req in sub.required_columns: resolved = graph.resolve_side_effect(req) if resolved not in known_columns: @@ -87,8 +100,23 @@ def create( f"Column '{name}' requires '{req}' (resolved to '{resolved}') which is not a known producer." ) if resolved == name: - continue # skip self-dependency + continue + if resolved not in resolved_required: + resolved_required.append(resolved) graph.add_edge(upstream=resolved, downstream=name) + graph.set_required_columns(name, resolved_required) + + if sub.skip is not None: + for skip_col in sub.skip.columns: + resolved = graph.resolve_side_effect(skip_col) + if resolved not in known_columns: + raise ValueError( + f"Column '{name}' skip.when references '{skip_col}' " + f"(resolved to '{resolved}') which is not a known producer." + ) + if resolved == name: + continue + graph.add_edge(upstream=resolved, downstream=name) # Validate acyclicity graph.get_topological_order() @@ -122,6 +150,19 @@ def set_side_effect(self, side_effect_col: str, producer: str) -> None: f"Use distinct side-effect column names for each pipeline stage." ) self._side_effect_map[side_effect_col] = producer + self._producer_to_side_effect_map.setdefault(producer, []).append(side_effect_col) + + def set_required_columns(self, column: str, required: list[str]) -> None: + """Store producer-resolved ``required_columns`` for skip propagation.""" + self._required_columns[column] = required + + def set_propagate_skip(self, column: str, propagate: bool) -> None: + """Store whether *column* should auto-skip when an upstream was skipped.""" + self._propagate_skip[column] = propagate + + def set_skip_config(self, column: str, skip_config: SkipConfig) -> None: + """Attach a ``SkipConfig`` to *column*.""" + self._skip_configs[column] = skip_config def resolve_side_effect(self, column: str) -> str: """Resolve a column name through the side-effect map. @@ -141,6 +182,22 @@ def get_downstream_columns(self, column: str) -> set[str]: """Columns that depend on *column*.""" return set(self._downstream.get(column, set())) + def get_required_columns(self, column: str) -> list[str]: + """Producer-resolved ``required_columns`` for *column* (data dependencies only).""" + return list(self._required_columns.get(column, [])) + + def get_skip_config(self, column: str) -> SkipConfig | None: + """Return the ``SkipConfig`` for *column*, or ``None`` if not configured.""" + return self._skip_configs.get(column) + + def should_propagate_skip(self, column: str) -> bool: + """Whether *column* auto-skips when an upstream was skipped.""" + return self._propagate_skip.get(column, True) + + def get_side_effect_columns(self, column: str) -> list[str]: + """Return side-effect column names produced by *column*.""" + return list(self._producer_to_side_effect_map.get(column, [])) + def get_strategy(self, column: str) -> GenerationStrategy: return self._strategies[column] diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py index 3adad1456..5ddbfec8c 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Callable import data_designer.lazy_heavy_imports as lazy +from data_designer.engine.dataset_builders.utils.skip_tracker import strip_skip_metadata_from_records if TYPE_CHECKING: import pandas as pd @@ -66,10 +67,10 @@ def has_row_group(self, row_group: int) -> bool: return row_group in self._buffers def get_dataframe(self, row_group: int) -> pd.DataFrame: - """Return the row group as a DataFrame (excluding dropped rows).""" + """Return the row group as a DataFrame (excluding dropped rows, stripping skip metadata).""" dropped = self._dropped.get(row_group, set()) rows = [row for i, row in enumerate(self._buffers[row_group]) if i not in dropped] - return lazy.pd.DataFrame(rows) + return lazy.pd.DataFrame(strip_skip_metadata_from_records(rows)) def replace_dataframe(self, row_group: int, df: pd.DataFrame) -> None: """Replace the buffer for a row group from a DataFrame (non-dropped rows only). diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_evaluator.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_evaluator.py new file mode 100644 index 000000000..3fcd1f7df --- /dev/null +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_evaluator.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Skip expression evaluation for conditional column generation.""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from typing import TYPE_CHECKING + +from jinja2 import StrictUndefined +from jinja2.exceptions import SecurityError, TemplateSyntaxError, UndefinedError +from jinja2.nativetypes import NativeEnvironment +from jinja2.sandbox import SandboxedEnvironment + +from data_designer.engine.dataset_builders.utils.skip_tracker import SKIPPED_COLUMNS_RECORD_KEY +from data_designer.engine.processing.utils import deserialize_json_values + +if TYPE_CHECKING: + from jinja2 import Template + +logger = logging.getLogger(__name__) + + +class NativeSandboxedEnvironment(SandboxedEnvironment, NativeEnvironment): + """Sandboxed environment that returns native Python types instead of strings. + + Uses ``StrictUndefined`` so that references to missing variables raise + ``UndefinedError`` instead of silently returning a truthy ``Undefined`` + object (which would cause every row to be skipped on a typo). + """ + + +_env = NativeSandboxedEnvironment(undefined=StrictUndefined) + + +def evaluate_skip_when(expression: str, record: dict) -> bool: + """Render *expression* against *record*; return ``True`` if result is truthy. + + The caller is responsible for passing a raw record dict — deserialization + of JSON string values is handled here so both sync and async engines get + identical behavior. On expected evaluation failures (``UndefinedError``, + ``SecurityError``, ``TemplateSyntaxError``, ``TypeError``, ``ValueError``) + a warning is logged and ``True`` is returned (fail-safe: skip the row + rather than making an expensive LLM call on a row with unknown filter + status). Unexpected exceptions propagate to the caller. + """ + try: + template = _compile_skip_template(expression) + deserialized = deserialize_json_values(record) + result = template.render(deserialized) + return bool(result) + except (UndefinedError, SecurityError, TemplateSyntaxError, TypeError, ValueError): + logger.warning( + "skip.when evaluation failed for expression %r; treating as truthy (cell will be skipped)", + expression, + exc_info=True, + ) + return True + + +def get_skipped_column_names(record: dict) -> set[str]: + """Return a *copy* of skipped producer column names for this row (empty if unset).""" + return set(record.get(SKIPPED_COLUMNS_RECORD_KEY, set())) + + +def should_skip_by_propagation( + required_columns: list[str], + skipped_columns_for_row: set[str], +) -> bool: + """Return ``True`` if any required column was skipped. + + The caller is responsible for checking ``propagate_skip`` on the column + config *before* calling this function (see ``ExecutionGraph.should_propagate_skip``). + """ + return not skipped_columns_for_row.isdisjoint(required_columns) + + +def should_skip_column_for_record( + record: dict, + *, + propagate_skip: bool, + required_columns: list[str], + skip_config_when: str | None, +) -> bool: + """Unified skip decision for a single cell/record. + + Shared by both the sync and async engines so the logic stays in sync. + Checks propagation first (cheaper), then the expression gate. + + Args: + record: Current row dict (may contain ``__internal_skipped_columns``). + propagate_skip: Whether this column auto-skips on upstream skips. + required_columns: Config-level data dependencies for *column*. + skip_config_when: The ``skip.when`` Jinja2 expression, or ``None``. + """ + skipped_cols = get_skipped_column_names(record) + if propagate_skip and should_skip_by_propagation(required_columns, skipped_cols): + return True + if skip_config_when is not None: + return evaluate_skip_when(skip_config_when, record) + return False + + +@lru_cache(maxsize=64) +def _compile_skip_template(expression: str) -> Template: + return _env.from_string(expression) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_tracker.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_tracker.py new file mode 100644 index 000000000..8529e0687 --- /dev/null +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_tracker.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Record-inline skip tracking for conditional column generation. + +All reads, writes, and DataFrame-stripping of the ``__internal_skipped_columns`` key go +through this module so sync, async, and buffer code do not diverge. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final + +SKIPPED_COLUMNS_RECORD_KEY: Final[str] = "__internal_skipped_columns" +SKIP_METADATA_RESTORE_ID_COLUMN_PREFIX: Final[str] = "__internal_skip_restore_id" + + +@dataclass(frozen=True, slots=True) +class SkipMetadataRestoreContext: + """Metadata needed to restore skip provenance after a DataFrame round-trip.""" + + restore_id_column: str + source_ids: set[str] + skipped_columns_by_source_id: dict[str, set[str]] + + +def apply_skip_to_record( + record: dict, + *, + column_name: str, + cell_value: bool | int | float | str | None, + side_effect_columns: Sequence[str], +) -> None: + """Mutate *record* in place: skip marker, primary cell value, side effects cleared. + + Side-effect columns (e.g. ``__trace``, ``__reasoning_content``) are set to + ``None`` because the generator never ran — without this, records would have + inconsistent keys, breaking DataFrame construction and leaving stale or + missing values visible to downstream columns. + """ + skipped: set[str] = record.setdefault(SKIPPED_COLUMNS_RECORD_KEY, set()) + skipped.add(column_name) + record[column_name] = cell_value + for se_col in side_effect_columns: + record[se_col] = None + skipped.add(se_col) + + +def strip_skip_metadata_for_dataframe_row(record: dict) -> dict: + """Shallow copy of *record* without skip metadata — safe for ``pd.DataFrame(rows)``.""" + return {k: v for k, v in record.items() if k != SKIPPED_COLUMNS_RECORD_KEY} + + +def strip_skip_metadata_from_records(records: Sequence[dict]) -> list[dict]: + """Map :func:`strip_skip_metadata_for_dataframe_row` over *records*.""" + return [strip_skip_metadata_for_dataframe_row(r) for r in records] + + +def prepare_records_for_skip_metadata_round_trip( + records: Sequence[dict], +) -> tuple[list[dict], SkipMetadataRestoreContext | None]: + """Prepare records for a DataFrame round-trip while preserving skip metadata. + + Returns stripped records ready for ``pd.DataFrame(...)``. If any record has + skip metadata, injects a hidden restore-ID column and returns a context that + can later be passed to :func:`restore_skip_metadata`. + """ + if not any(SKIPPED_COLUMNS_RECORD_KEY in record for record in records): + return strip_skip_metadata_from_records(records), None + + restore_id_column = _choose_restore_id_column(records) + prepared_records: list[dict] = [] + source_ids: set[str] = set() + skipped_columns_by_source_id: dict[str, set[str]] = {} + + for index, record in enumerate(records): + source_id = str(index) + source_ids.add(source_id) + prepared_record = strip_skip_metadata_for_dataframe_row(record) + prepared_record[restore_id_column] = source_id + prepared_records.append(prepared_record) + + meta = record.get(SKIPPED_COLUMNS_RECORD_KEY) + if meta is not None: + skipped_columns_by_source_id[source_id] = set(meta) + + return prepared_records, SkipMetadataRestoreContext( + restore_id_column=restore_id_column, + source_ids=source_ids, + skipped_columns_by_source_id=skipped_columns_by_source_id, + ) + + +def restore_skip_metadata( + records: Sequence[dict], + *, + context: SkipMetadataRestoreContext, + allow_resize: bool, +) -> None: + """Restore skip provenance using hidden restore IDs instead of row position.""" + restored_source_ids: list[str] = [] + for record in records: + if context.restore_id_column not in record: + raise ValueError( + f"Records returned from the DataFrame round-trip must preserve " + f"the internal column {context.restore_id_column!r} so skip " + "provenance can be restored." + ) + + source_id = str(record.pop(context.restore_id_column)) + if source_id not in context.source_ids: + raise ValueError( + f"Record returned unknown restore ID {source_id!r}. Skip provenance " + "can only be restored for rows derived from the original input." + ) + + restored_source_ids.append(source_id) + meta = context.skipped_columns_by_source_id.get(source_id) + if meta is not None: + record[SKIPPED_COLUMNS_RECORD_KEY] = set(meta) + + if not allow_resize: + if len(restored_source_ids) != len(context.source_ids) or set(restored_source_ids) != context.source_ids: + raise ValueError( + "Full-column generation changed the row identity mapping while " + "allow_resize=False. Returned rows must preserve a 1:1 mapping " + "to the original input so skip provenance can be restored." + ) + + +def _choose_restore_id_column(records: Sequence[dict]) -> str: + candidate = SKIP_METADATA_RESTORE_ID_COLUMN_PREFIX + suffix = 0 + while any(candidate in record for record in records): + suffix += 1 + candidate = f"{SKIP_METADATA_RESTORE_ID_COLUMN_PREFIX}_{suffix}" + return candidate diff --git a/packages/data-designer-engine/src/data_designer/engine/validation.py b/packages/data-designer-engine/src/data_designer/engine/validation.py index bc8a792a1..600fde233 100644 --- a/packages/data-designer-engine/src/data_designer/engine/validation.py +++ b/packages/data-designer-engine/src/data_designer/engine/validation.py @@ -39,6 +39,9 @@ class ViolationType(str, Enum): INVALID_MODEL_CONFIG = "invalid_model_config" INVALID_REFERENCE = "invalid_reference" PROMPT_WITHOUT_REFERENCES = "prompt_without_references" + SKIP_REFERENCE_MISSING = "skip_reference_missing" + SKIP_ON_SAMPLER_SEED = "skip_on_sampler_seed" + SKIP_WITH_ALLOW_RESIZE = "skip_with_allow_resize" class ViolationLevel(str, Enum): @@ -66,6 +69,7 @@ def validate_data_designer_config( violations.extend(validate_prompt_templates(columns=columns, allowed_references=allowed_references)) violations.extend(validate_code_validation(columns=columns)) violations.extend(validate_expression_references(columns=columns, allowed_references=allowed_references)) + violations.extend(validate_skip_references(columns=columns, allowed_references=allowed_references)) violations.extend(validate_columns_not_all_dropped(columns=columns)) violations.extend(validate_drop_columns_processor(columns=columns, processor_configs=processor_configs)) violations.extend(validate_schema_transform_processor(columns=columns, processor_configs=processor_configs)) @@ -378,6 +382,57 @@ def validate_local_only_columns( return violations +def validate_skip_references( + columns: list[ColumnConfigT], + allowed_references: list[str], +) -> list[Violation]: + """Validate ``skip.when`` expressions: reference existence, type scope, and ``allow_resize`` conflicts.""" + violations: list[Violation] = [] + for column in columns: + if column.skip is None: + continue + + if column.column_type in ("sampler", "seed-dataset"): + violations.append( + Violation( + column=column.name, + type=ViolationType.SKIP_ON_SAMPLER_SEED, + message=( + f"skip is not supported on {column.column_type} columns. " + "Sampler/seed columns are collapsed into shared multi-column generators " + "and cannot be skipped individually." + ), + level=ViolationLevel.ERROR, + ) + ) + + if getattr(column, "allow_resize", False): + violations.append( + Violation( + column=column.name, + type=ViolationType.SKIP_WITH_ALLOW_RESIZE, + message="skip and allow_resize cannot be used together on the same column.", + level=ViolationLevel.ERROR, + ) + ) + + for ref in column.skip.columns: + if ref not in allowed_references: + violations.append( + Violation( + column=column.name, + type=ViolationType.SKIP_REFERENCE_MISSING, + message=( + f"skip.when expression for column '{column.name}' references " + f"'{ref}' which is not a known column." + ), + level=ViolationLevel.ERROR, + ) + ) + + return violations + + def _get_string_formatter_references(template: str, allowed_references: list[str]) -> list[str]: return [ k[1].strip() diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py index 32a9c69ce..cea3e4ccf 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py @@ -10,6 +10,7 @@ import pytest import data_designer.lazy_heavy_imports as lazy +from data_designer.config.base import SkipConfig from data_designer.config.column_configs import ( CustomColumnConfig, ExpressionColumnConfig, @@ -1596,3 +1597,244 @@ async def test_scheduler_downstream_interleaves_with_upstream() -> None: f"First judge dispatched at {first_judge_dispatched:.4f}, " f"last gen dispatched at {last_gen_dispatched:.4f}." ) + + +# -- Skip / conditional generation tests (async engine) ----------------------- + + +@pytest.mark.asyncio(loop_scope="session") +async def test_scheduler_skip_cell_by_cell_with_propagation() -> None: + """Cell-by-cell column skips rows via expression gate, downstream propagates. + + Pipeline: seed(sampler) -> review(cell, skip.when seed<2) -> complaint(cell, propagate_skip) + Rows with seed < 2 should be skipped for review and propagated to complaint. + """ + provider = _mock_provider() + num_records = 4 + + configs = [ + SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}), + LLMTextColumnConfig( + name="review", + prompt="{{ seed }}", + model_alias=MODEL_ALIAS, + skip=SkipConfig(when="{{ seed < 2 }}"), + ), + LLMTextColumnConfig( + name="complaint", + prompt="{{ review }}", + model_alias=MODEL_ALIAS, + propagate_skip=True, + ), + ] + strategies = { + "seed": GenerationStrategy.FULL_COLUMN, + "review": GenerationStrategy.CELL_BY_CELL, + "complaint": GenerationStrategy.CELL_BY_CELL, + } + + class IntSeedGenerator(FromScratchColumnGenerator[ExpressionColumnConfig]): + @staticmethod + def get_generation_strategy() -> GenerationStrategy: + return GenerationStrategy.FULL_COLUMN + + def generate(self, data: lazy.pd.DataFrame) -> lazy.pd.DataFrame: + return data + + def generate_from_scratch(self, num_records: int) -> lazy.pd.DataFrame: + return lazy.pd.DataFrame({"seed": list(range(num_records))}) + + generators: dict[str, ColumnGenerator] = { + "seed": IntSeedGenerator(config=_expr_config("seed"), resource_provider=provider), + "review": MockCellGenerator(config=_expr_config("review"), resource_provider=provider), + "complaint": MockCellGenerator(config=_expr_config("complaint"), resource_provider=provider), + } + + storage = MagicMock() + storage.dataset_name = "test" + storage.get_file_paths.return_value = {} + buffer_mgr = RowGroupBufferManager(storage) + + graph = ExecutionGraph.create(configs, strategies) + row_groups = [(0, num_records)] + tracker = CompletionTracker.with_graph(graph, row_groups) + + scheduler = AsyncTaskScheduler( + generators=generators, + graph=graph, + tracker=tracker, + row_groups=row_groups, + buffer_manager=buffer_mgr, + trace=True, + num_records=num_records, + buffer_size=num_records, + ) + await asyncio.wait_for(scheduler.run(), timeout=10.0) + + assert tracker.is_row_group_complete(0, num_records, ["seed", "review", "complaint"]) + + for ri in range(num_records): + row = buffer_mgr.get_row(0, ri) + seed_val = row["seed"] + if seed_val < 2: + assert row.get("review") is None, f"row {ri}: review should be skipped (seed={seed_val})" + assert row.get("complaint") is None, f"row {ri}: complaint should propagate skip (seed={seed_val})" + else: + assert row.get("review") is not None, f"row {ri}: review should be generated (seed={seed_val})" + assert row.get("complaint") is not None, f"row {ri}: complaint should be generated (seed={seed_val})" + + +@pytest.mark.asyncio(loop_scope="session") +async def test_scheduler_skip_propagates_through_side_effect_dependency() -> None: + """A downstream dependency on a skipped side-effect should auto-skip. + + Pipeline: seed(sampler) -> review(cell, skip.when seed<2, produces + review__trace) -> complaint(cell, depends on review__trace, + propagate_skip=True). + """ + provider = _mock_provider() + num_records = 4 + + configs = [ + SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}), + LLMTextColumnConfig( + name="review", + prompt="{{ seed }}", + model_alias=MODEL_ALIAS, + with_trace="last_message", + skip=SkipConfig(when="{{ seed < 2 }}"), + ), + LLMTextColumnConfig( + name="complaint", + prompt="{{ review__trace }}", + model_alias=MODEL_ALIAS, + propagate_skip=True, + ), + ] + strategies = { + "seed": GenerationStrategy.FULL_COLUMN, + "review": GenerationStrategy.CELL_BY_CELL, + "complaint": GenerationStrategy.CELL_BY_CELL, + } + + class IntSeedGenerator(FromScratchColumnGenerator[ExpressionColumnConfig]): + @staticmethod + def get_generation_strategy() -> GenerationStrategy: + return GenerationStrategy.FULL_COLUMN + + def generate(self, data: lazy.pd.DataFrame) -> lazy.pd.DataFrame: + return data + + def generate_from_scratch(self, num_records: int) -> lazy.pd.DataFrame: + return lazy.pd.DataFrame({"seed": list(range(num_records))}) + + generators: dict[str, ColumnGenerator] = { + "seed": IntSeedGenerator(config=_expr_config("seed"), resource_provider=provider), + "review": MockCellGenerator(config=_expr_config("review"), resource_provider=provider), + "complaint": MockCellGenerator(config=_expr_config("complaint"), resource_provider=provider), + } + + storage = MagicMock() + storage.dataset_name = "test" + storage.get_file_paths.return_value = {} + buffer_mgr = RowGroupBufferManager(storage) + + graph = ExecutionGraph.create(configs, strategies) + row_groups = [(0, num_records)] + tracker = CompletionTracker.with_graph(graph, row_groups) + + scheduler = AsyncTaskScheduler( + generators=generators, + graph=graph, + tracker=tracker, + row_groups=row_groups, + buffer_manager=buffer_mgr, + trace=True, + num_records=num_records, + buffer_size=num_records, + ) + await asyncio.wait_for(scheduler.run(), timeout=10.0) + + assert tracker.is_row_group_complete(0, num_records, ["seed", "review", "complaint"]) + + for ri in range(num_records): + row = buffer_mgr.get_row(0, ri) + seed_val = row["seed"] + if seed_val < 2: + assert row.get("review") is None, f"row {ri}: review should be skipped (seed={seed_val})" + assert row.get("review__trace") is None, f"row {ri}: review__trace should be cleared on skip" + assert row.get("complaint") is None, f"row {ri}: complaint should propagate skip (seed={seed_val})" + else: + assert row.get("complaint") is not None, f"row {ri}: complaint should be generated (seed={seed_val})" + + +@pytest.mark.asyncio(loop_scope="session") +async def test_scheduler_skip_full_column_batch() -> None: + """Full-column (batch) generator skips rows via expression gate. + + Pipeline: seed(sampler) -> review(full_column, skip.when seed<2) + Only active (non-skipped) rows should be passed to the generator. + """ + provider = _mock_provider() + num_records = 4 + + configs = [ + SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}), + LLMTextColumnConfig( + name="review", + prompt="{{ seed }}", + model_alias=MODEL_ALIAS, + skip=SkipConfig(when="{{ seed < 2 }}"), + ), + ] + strategies = { + "seed": GenerationStrategy.FULL_COLUMN, + "review": GenerationStrategy.FULL_COLUMN, + } + + class IntSeedGenerator(FromScratchColumnGenerator[ExpressionColumnConfig]): + @staticmethod + def get_generation_strategy() -> GenerationStrategy: + return GenerationStrategy.FULL_COLUMN + + def generate(self, data: lazy.pd.DataFrame) -> lazy.pd.DataFrame: + return data + + def generate_from_scratch(self, num_records: int) -> lazy.pd.DataFrame: + return lazy.pd.DataFrame({"seed": list(range(num_records))}) + + generators: dict[str, ColumnGenerator] = { + "seed": IntSeedGenerator(config=_expr_config("seed"), resource_provider=provider), + "review": MockFullColumnGenerator(config=_expr_config("review"), resource_provider=provider), + } + + storage = MagicMock() + storage.dataset_name = "test" + storage.get_file_paths.return_value = {} + buffer_mgr = RowGroupBufferManager(storage) + + graph = ExecutionGraph.create(configs, strategies) + row_groups = [(0, num_records)] + tracker = CompletionTracker.with_graph(graph, row_groups) + + scheduler = AsyncTaskScheduler( + generators=generators, + graph=graph, + tracker=tracker, + row_groups=row_groups, + buffer_manager=buffer_mgr, + trace=True, + num_records=num_records, + buffer_size=num_records, + ) + await asyncio.wait_for(scheduler.run(), timeout=10.0) + + assert tracker.is_row_group_complete(0, num_records, ["seed", "review"]) + + for ri in range(num_records): + row = buffer_mgr.get_row(0, ri) + seed_val = row["seed"] + if seed_val < 2: + assert row.get("review") is None, f"row {ri}: review should be skipped (seed={seed_val})" + else: + assert row["review"] == "batch_val", f"row {ri}: review should be generated (seed={seed_val})" diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py index 477978a38..dd72b8461 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py @@ -11,6 +11,7 @@ import data_designer.engine.dataset_builders.dataset_builder as builder_mod import data_designer.lazy_heavy_imports as lazy +from data_designer.config.base import SkipConfig from data_designer.config.column_configs import CustomColumnConfig, LLMTextColumnConfig, SamplerColumnConfig from data_designer.config.config_builder import DataDesignerConfigBuilder from data_designer.config.custom_column import custom_column_generator @@ -502,12 +503,17 @@ def test_fan_out_with_threads_uses_early_shutdown_settings_from_resource_provide def test_fan_out_with_threads_passes_column_name_in_context( mock_executor_class: Mock, stub_resource_provider: Mock, - stub_test_config_builder: DataDesignerConfigBuilder, + stub_model_configs: dict[str, object], ) -> None: + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.add_column( + SamplerColumnConfig(name="some_id", sampler_type=SamplerType.UUID, params=UUIDSamplerParams()) + ) builder = DatasetBuilder( - data_designer_config=stub_test_config_builder.build(), + data_designer_config=config_builder.build(), resource_provider=stub_resource_provider, ) + builder.build_preview(num_records=1) mock_executor = Mock() mock_executor_class.return_value.__enter__ = Mock(return_value=mock_executor) @@ -537,12 +543,17 @@ def test_fan_out_with_threads_passes_column_name_in_context( def test_fan_out_with_async_passes_column_name_in_context( mock_executor_class: Mock, stub_resource_provider: Mock, - stub_test_config_builder: DataDesignerConfigBuilder, + stub_model_configs: dict[str, object], ) -> None: + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.add_column( + SamplerColumnConfig(name="some_id", sampler_type=SamplerType.UUID, params=UUIDSamplerParams()) + ) builder = DatasetBuilder( - data_designer_config=stub_test_config_builder.build(), + data_designer_config=config_builder.build(), resource_provider=stub_resource_provider, ) + builder.build_preview(num_records=1) mock_executor = Mock() @@ -794,6 +805,12 @@ def _resize_full_keep_first(df: pd.DataFrame) -> pd.DataFrame: return df.drop_duplicates(subset="seed_id").assign(filtered=True) +@custom_column_generator(required_columns=["seed_id"]) +def _resize_full_drop_seed_one(df: pd.DataFrame) -> pd.DataFrame: + """FULL_COLUMN: drop the row with seed_id == 1.""" + return df[df["seed_id"] != 1].reset_index(drop=True).assign(filtered=True) + + @custom_column_generator(required_columns=["seed_id"]) def _resize_cell_expand(row: dict) -> list[dict]: """CELL_BY_CELL: one row -> two rows (doubled).""" @@ -937,3 +954,445 @@ def test_allow_resize_multiple_batches( else: df = lazy.pd.read_parquet(final_path) assert len(df) == expected_total_rows + + +# skip metadata preservation tests + + +def _make_label_generator(label: str, *required: str): + """FULL_COLUMN generator that adds a column with a constant label value.""" + + @custom_column_generator(required_columns=list(required)) + def fn(df: pd.DataFrame) -> pd.DataFrame: + return df.assign(**{label: f"generated_{label}"}) + + return fn + + +def _make_label_generator_with_side_effect(label: str, side_effect_label: str, *required: str): + """FULL_COLUMN generator that adds a column plus one side-effect column.""" + + @custom_column_generator(required_columns=list(required), side_effect_columns=[side_effect_label]) + def fn(df: pd.DataFrame) -> pd.DataFrame: + return df.assign( + **{ + label: f"generated_{label}", + side_effect_label: f"generated_{side_effect_label}", + } + ) + + return fn + + +def test_skip_metadata_preserved_across_non_skip_aware_full_column( + stub_resource_provider, stub_model_configs, seed_data_setup +): + """Skip metadata must survive when a non-skip-aware FULL_COLUMN column runs + between a skip-setting column and a downstream propagating column. + + Scenario: rating(seed) -> review(skip.when) -> summary(no skip) -> complaint(propagate_skip) + Before the fix, summary's replace_buffer erased __internal_skipped_columns, + causing complaint to generate for rows that should have been skipped. + """ + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator("review", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="summary", + generator_function=_make_label_generator("summary", "text"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=False, + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="complaint", + generator_function=_make_label_generator("complaint", "review"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + skipped_ids = {1, 2} + for _, row in result.iterrows(): + if row["seed_id"] in skipped_ids: + assert row["review"] is None or lazy.pd.isna(row["review"]), ( + f"seed_id={row['seed_id']}: review should be skipped" + ) + assert row["complaint"] is None or lazy.pd.isna(row["complaint"]), ( + f"seed_id={row['seed_id']}: complaint should propagate skip from review" + ) + else: + assert row["complaint"] == "generated_complaint", f"seed_id={row['seed_id']}: complaint should be generated" + + +def test_skip_metadata_preserved_when_no_rows_skipped_for_current_column( + stub_resource_provider, stub_model_configs, seed_data_setup +): + """The has_skipped=False fallthrough must preserve sibling skip metadata. + + Scenario: review(skip.when seed_id<3) -> analysis(propagate_skip, required_columns=[review]) + analysis can_skip=True (via propagation) but no rows are skipped by analysis's + own expression (it has none). The has_skipped=False fallthrough must still + preserve review's skip metadata so propagation works. + """ + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator("review", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="analysis", + generator_function=_make_label_generator("analysis", "review"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + skipped_ids = {1, 2} + for _, row in result.iterrows(): + if row["seed_id"] in skipped_ids: + assert row["analysis"] is None or lazy.pd.isna(row["analysis"]), ( + f"seed_id={row['seed_id']}: analysis should propagate skip from review" + ) + else: + assert row["analysis"] == "generated_analysis", f"seed_id={row['seed_id']}: analysis should be generated" + + +def test_skip_propagation_resolves_side_effect_dependencies_in_sync_builder( + stub_resource_provider, stub_model_configs, seed_data_setup +): + """A downstream dependency on a skipped side-effect should auto-skip. + + Scenario: review(skip.when, produces review_side_effect) -> + analysis(required_columns=[review_side_effect], propagate_skip=True). + """ + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator_with_side_effect("review", "review_side_effect", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="analysis", + generator_function=_make_label_generator("analysis", "review_side_effect"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + skipped_ids = {1, 2} + for _, row in result.iterrows(): + if row["seed_id"] in skipped_ids: + assert row["review_side_effect"] is None or lazy.pd.isna(row["review_side_effect"]), ( + f"seed_id={row['seed_id']}: review_side_effect should be cleared when review is skipped" + ) + assert row["analysis"] is None or lazy.pd.isna(row["analysis"]), ( + f"seed_id={row['seed_id']}: analysis should propagate skip from review" + ) + else: + assert row["analysis"] == "generated_analysis", f"seed_id={row['seed_id']}: analysis should be generated" + + +def test_skip_metadata_restore_preserves_row_identity_across_allow_resize_full_column( + stub_resource_provider, stub_model_configs, seed_data_setup +): + """Filtering out a skipped row must not transfer its skip provenance to surviving rows.""" + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator("review", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id == 1 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="filtered", + generator_function=_resize_full_drop_seed_one, + generation_strategy=GenerationStrategy.FULL_COLUMN, + allow_resize=True, + propagate_skip=False, + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="analysis", + generator_function=_make_label_generator("analysis", "review"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + assert result["seed_id"].tolist() == [2, 3, 4, 5] + assert result["analysis"].tolist() == ["generated_analysis"] * 4 + + +def test_allow_resize_column_not_blocked_by_upstream_skip(stub_resource_provider, stub_model_configs, seed_data_setup): + """An allow_resize=True column depending on a skippable upstream must not + enter the skip-aware branch (which enforces 1:1 row counts). + + Before the fix, _column_can_skip returned True for allow_resize columns + with propagate_skip=True and required_columns pointing to a skippable + upstream, causing a DatasetGenerationError on the row-count check. + """ + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator("review", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="expanded", + generator_function=_make_resize_full_expand(2, "expanded", "copy"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + allow_resize=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + assert len(result) == 10 + + +def test_skip_chained_transitive_propagation_through_three_levels( + stub_resource_provider, stub_model_configs, seed_data_setup +) -> None: + """Skip at level 1 must propagate transitively through levels 2, 3, and 4. + + Pipeline: seed_id(seed) -> L1(skip.when) -> L2(propagate) -> L3(propagate) -> L4(propagate) + """ + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="L1", + generator_function=_make_label_generator("L1", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="L2", + generator_function=_make_label_generator("L2", "L1"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="L3", + generator_function=_make_label_generator("L3", "L2"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="L4", + generator_function=_make_label_generator("L4", "L3"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + assert len(result) == 5 + skipped_ids = {1, 2} + for _, row in result.iterrows(): + if row["seed_id"] in skipped_ids: + for col in ("L1", "L2", "L3", "L4"): + assert row[col] is None or lazy.pd.isna(row[col]), ( + f"seed_id={row['seed_id']}: {col} should be skipped transitively" + ) + else: + for col in ("L1", "L2", "L3", "L4"): + assert row[col] == f"generated_{col}", f"seed_id={row['seed_id']}: {col} should be generated" + + +def test_skip_two_independent_gates_in_same_pipeline( + stub_resource_provider, stub_model_configs, seed_data_setup +) -> None: + """Two columns with independent skip.when expressions; downstream propagates from both. + + Pipeline: seed_id(seed) -> gate_a(skip seed_id<3) -> gate_b(skip seed_id>4) -> merge(propagate) + merge should be skipped when *either* gate_a or gate_b was skipped. + """ + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="gate_a", + generator_function=_make_label_generator("gate_a", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="gate_b", + generator_function=_make_label_generator("gate_b", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id > 4 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="merge", + generator_function=_make_label_generator("merge", "gate_a", "gate_b"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + assert len(result) == 5 + for _, row in result.iterrows(): + sid = row["seed_id"] + if sid < 3 or sid > 4: + assert row["merge"] is None or lazy.pd.isna(row["merge"]), ( + f"seed_id={sid}: merge should be skipped (gate_a or gate_b skipped)" + ) + else: + assert row["merge"] == "generated_merge", f"seed_id={sid}: merge should be generated" + + +def test_skip_custom_value_preserved_in_output(stub_resource_provider, stub_model_configs, seed_data_setup) -> None: + """Custom skip.value should appear in the final DataFrame instead of None.""" + sentinel = "__SKIPPED__" + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator("review", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}", value=sentinel), + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + assert len(result) == 5 + skipped_ids = {1, 2} + for _, row in result.iterrows(): + if row["seed_id"] in skipped_ids: + assert row["review"] == sentinel, f"seed_id={row['seed_id']}: review should have custom skip value" + else: + assert row["review"] == "generated_review", f"seed_id={row['seed_id']}: review should be generated" + + +def test_skip_row_count_preserved_across_pipeline(stub_resource_provider, stub_model_configs, seed_data_setup) -> None: + """Skip must never change the row count — all 5 seed rows must survive.""" + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"]))) + + config_builder.add_column( + CustomColumnConfig( + name="review", + generator_function=_make_label_generator("review", "seed_id"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + skip=SkipConfig(when="{{ seed_id < 3 }}"), + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="analysis", + generator_function=_make_label_generator("analysis", "review"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + config_builder.add_column( + CustomColumnConfig( + name="summary", + generator_function=_make_label_generator("summary", "analysis"), + generation_strategy=GenerationStrategy.FULL_COLUMN, + propagate_skip=True, + ) + ) + + builder = DatasetBuilder( + data_designer_config=config_builder.build(), + resource_provider=stub_resource_provider, + ) + result = builder.build_preview(num_records=5) + + assert len(result) == 5, "Skip must not change the row count" + assert result["seed_id"].tolist() == [1, 2, 3, 4, 5] diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_execution_graph.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_execution_graph.py index c1aa3eeef..dfd219fd5 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_execution_graph.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_execution_graph.py @@ -5,6 +5,7 @@ import pytest +from data_designer.config.base import SkipConfig from data_designer.config.column_configs import ( ExpressionColumnConfig, GenerationStrategy, @@ -112,6 +113,7 @@ def test_side_effect_column_resolution() -> None: assert graph.get_upstream_columns("trace_len") == {"summary"} assert graph.get_downstream_columns("summary") == {"trace_len"} + assert graph.get_required_columns("trace_len") == ["summary"] def test_reasoning_content_side_effect() -> None: @@ -459,3 +461,104 @@ def test_judge_column_dependency() -> None: graph = ExecutionGraph.create(configs, strategies) assert graph.get_upstream_columns("judge") == {"text"} + + +# -- Skip metadata accessors ------------------------------------------------ + + +def _build_skip_pipeline_graph() -> ExecutionGraph: + """gate(sampler) -> review(skip.when, with_trace) -> analysis(propagate) -> summary(no propagate).""" + configs = [ + SamplerColumnConfig(name="gate", sampler_type=SamplerType.CATEGORY, params={"values": [0, 1]}), + LLMTextColumnConfig( + name="review", + prompt="{{ gate }}", + model_alias=MODEL_ALIAS, + with_trace="last_message", + skip=SkipConfig(when="{{ gate == 0 }}"), + ), + LLMTextColumnConfig( + name="analysis", + prompt="{{ review }}", + model_alias=MODEL_ALIAS, + propagate_skip=True, + ), + LLMTextColumnConfig( + name="summary", + prompt="{{ analysis }}", + model_alias=MODEL_ALIAS, + propagate_skip=False, + ), + ] + strategies = { + "gate": GenerationStrategy.FULL_COLUMN, + "review": GenerationStrategy.CELL_BY_CELL, + "analysis": GenerationStrategy.CELL_BY_CELL, + "summary": GenerationStrategy.CELL_BY_CELL, + } + return ExecutionGraph.create(configs, strategies) + + +def test_skip_config_returned_for_gated_column() -> None: + graph = _build_skip_pipeline_graph() + skip_cfg = graph.get_skip_config("review") + assert skip_cfg is not None + assert skip_cfg.when == "{{ gate == 0 }}" + + +def test_skip_config_returns_none_for_ungated_column() -> None: + graph = _build_skip_pipeline_graph() + assert graph.get_skip_config("gate") is None + assert graph.get_skip_config("analysis") is None + + +def test_should_propagate_skip_explicit_values() -> None: + graph = _build_skip_pipeline_graph() + assert graph.should_propagate_skip("analysis") is True + assert graph.should_propagate_skip("summary") is False + + +def test_should_propagate_skip_defaults_true() -> None: + graph = _build_skip_pipeline_graph() + assert graph.should_propagate_skip("gate") is True + assert graph.should_propagate_skip("review") is True + + +def test_get_required_columns_for_skip_pipeline() -> None: + graph = _build_skip_pipeline_graph() + assert graph.get_required_columns("review") == ["gate"] + assert graph.get_required_columns("analysis") == ["review"] + assert graph.get_required_columns("summary") == ["analysis"] + + +def test_get_side_effect_columns_for_skip_pipeline() -> None: + graph = _build_skip_pipeline_graph() + assert graph.get_side_effect_columns("review") == ["review__trace"] + assert graph.get_side_effect_columns("analysis") == [] + + +def test_side_effect_dependency_resolves_to_producer() -> None: + graph = _build_skip_pipeline_graph() + assert graph.resolve_side_effect("review__trace") == "review" + + +def test_skip_when_columns_create_dag_edges() -> None: + """skip.when referencing a column should create an edge in the DAG.""" + configs = [ + SamplerColumnConfig(name="gate", sampler_type=SamplerType.CATEGORY, params={"values": [0, 1]}), + SamplerColumnConfig(name="data", sampler_type=SamplerType.CATEGORY, params={"values": ["x"]}), + LLMTextColumnConfig( + name="output", + prompt="{{ data }}", + model_alias=MODEL_ALIAS, + skip=SkipConfig(when="{{ gate == 0 }}"), + ), + ] + strategies = { + "gate": GenerationStrategy.FULL_COLUMN, + "data": GenerationStrategy.FULL_COLUMN, + "output": GenerationStrategy.CELL_BY_CELL, + } + graph = ExecutionGraph.create(configs, strategies) + assert "gate" in graph.get_upstream_columns("output") + assert "data" in graph.get_upstream_columns("output") diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_evaluator.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_evaluator.py new file mode 100644 index 000000000..5eb05e464 --- /dev/null +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_evaluator.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from jinja2 import StrictUndefined + +from data_designer.engine.dataset_builders.utils.skip_evaluator import ( + NativeSandboxedEnvironment, + evaluate_skip_when, + should_skip_by_propagation, + should_skip_column_for_record, +) +from data_designer.engine.dataset_builders.utils.skip_tracker import SKIPPED_COLUMNS_RECORD_KEY + + +def test_native_sandboxed_environment_returns_native_types() -> None: + env = NativeSandboxedEnvironment(undefined=StrictUndefined) + result = env.from_string("{{ 1 + 1 }}").render() + assert result == 2 + assert type(result) is int + + +@pytest.mark.parametrize( + ("expression", "record", "expected"), + [ + pytest.param("{{ x == 0 }}", {"x": 0}, True, id="truthy-match"), + pytest.param("{{ x == 0 }}", {"x": 1}, False, id="falsy-no-match"), + pytest.param("{{ x }}", {"x": False}, False, id="native-false"), + pytest.param("{{ x }}", {"x": None}, False, id="native-none"), + pytest.param("{{ x }}", {"x": 0}, False, id="native-zero"), + pytest.param("{{ x }}", {"x": ""}, False, id="native-empty-string"), + pytest.param('{{ x.key == "val" }}', {"x": '{"key": "val"}'}, True, id="deserializes-json"), + ], +) +def test_evaluate_skip_when(expression: str, record: dict, expected: bool) -> None: + assert evaluate_skip_when(expression, record) is expected + + +def test_evaluate_skip_when_strict_undefined_returns_true() -> None: + """Missing variables trigger fail-safe: returns True (skip the row) and logs a warning.""" + assert evaluate_skip_when("{{ missing_var }}", {}) is True + + +@pytest.mark.parametrize( + ("required", "skipped", "expected"), + [ + pytest.param(["a", "b"], {"a"}, True, id="overlap"), + pytest.param(["a"], {"b"}, False, id="no-overlap"), + pytest.param([], {"a"}, False, id="empty-required"), + pytest.param(["a"], set(), False, id="empty-skipped"), + ], +) +def test_should_skip_by_propagation(required: list[str], skipped: set[str], expected: bool) -> None: + assert should_skip_by_propagation(required, skipped) is expected + + +# -- should_skip_column_for_record (unified decision) ----------------------- + +_UPSTREAM_SKIPPED_RECORD: dict = {SKIPPED_COLUMNS_RECORD_KEY: {"upstream_col"}, "upstream_col": None} + + +@pytest.mark.parametrize( + ("record", "propagate_skip", "required_columns", "skip_config_when", "expected"), + [ + pytest.param( + {"x": 1}, + True, + [], + None, + False, + id="no-gate-no-propagation", + ), + pytest.param( + _UPSTREAM_SKIPPED_RECORD, + True, + ["upstream_col"], + None, + True, + id="propagation-triggers-on-upstream-skip", + ), + pytest.param( + _UPSTREAM_SKIPPED_RECORD, + False, + ["upstream_col"], + None, + False, + id="propagation-disabled-ignores-upstream-skip", + ), + pytest.param( + {"gate": 0}, + True, + [], + "{{ gate == 0 }}", + True, + id="expression-truthy-skips", + ), + pytest.param( + {"gate": 1}, + True, + [], + "{{ gate == 0 }}", + False, + id="expression-falsy-does-not-skip", + ), + pytest.param( + {SKIPPED_COLUMNS_RECORD_KEY: {"dep"}, "dep": None, "gate": 999}, + True, + ["dep"], + "{{ gate == 0 }}", + True, + id="propagation-short-circuits-before-expression", + ), + pytest.param( + {SKIPPED_COLUMNS_RECORD_KEY: {"other"}, "gate": 0}, + True, + ["dep"], + "{{ gate == 0 }}", + True, + id="expression-evaluated-when-propagation-does-not-trigger", + ), + pytest.param( + {"gate": 1}, + True, + ["dep"], + "{{ gate == 0 }}", + False, + id="both-propagation-and-expression-false", + ), + ], +) +def test_should_skip_column_for_record( + record: dict, + propagate_skip: bool, + required_columns: list[str], + skip_config_when: str | None, + expected: bool, +) -> None: + assert ( + should_skip_column_for_record( + record, + propagate_skip=propagate_skip, + required_columns=required_columns, + skip_config_when=skip_config_when, + ) + is expected + ) diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_tracker.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_tracker.py new file mode 100644 index 000000000..7eeb68789 --- /dev/null +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_tracker.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from data_designer.engine.dataset_builders.utils.skip_evaluator import get_skipped_column_names +from data_designer.engine.dataset_builders.utils.skip_tracker import ( + SKIPPED_COLUMNS_RECORD_KEY, + apply_skip_to_record, + prepare_records_for_skip_metadata_round_trip, + restore_skip_metadata, + strip_skip_metadata_for_dataframe_row, + strip_skip_metadata_from_records, +) + + +def test_skipped_columns_record_key_value() -> None: + assert SKIPPED_COLUMNS_RECORD_KEY == "__internal_skipped_columns" + + +@pytest.mark.parametrize( + ("record", "expected"), + [ + pytest.param({}, set(), id="empty"), + pytest.param({SKIPPED_COLUMNS_RECORD_KEY: {"a", "b"}}, {"a", "b"}, id="populated"), + ], +) +def test_get_skipped_column_names(record: dict, expected: set[str]) -> None: + assert get_skipped_column_names(record) == expected + + +def test_get_skipped_column_names_returns_copy() -> None: + inner: set[str] = {"x"} + record = {SKIPPED_COLUMNS_RECORD_KEY: inner} + names = get_skipped_column_names(record) + names.add("y") + assert record[SKIPPED_COLUMNS_RECORD_KEY] == {"x"} + assert names == {"x", "y"} + + +def test_apply_skip_to_record_adds_skip_marker() -> None: + record: dict = {} + apply_skip_to_record( + record, + column_name="primary", + cell_value=None, + side_effect_columns=(), + ) + assert record[SKIPPED_COLUMNS_RECORD_KEY] == {"primary"} + + +@pytest.mark.parametrize( + "cell_value", + [None, True, False, 0, 42, 3.14, "skipped"], +) +def test_apply_skip_to_record_sets_cell_value(cell_value: bool | int | float | str | None) -> None: + record: dict = {} + apply_skip_to_record( + record, + column_name="col_a", + cell_value=cell_value, + side_effect_columns=(), + ) + assert record["col_a"] == cell_value + + +def test_apply_skip_to_record_clears_side_effects() -> None: + record: dict = {"se1": "keep-me", "se2": 99} + apply_skip_to_record( + record, + column_name="primary", + cell_value="pv", + side_effect_columns=("se1", "se2"), + ) + assert record["se1"] is None + assert record["se2"] is None + assert record["primary"] == "pv" + assert record[SKIPPED_COLUMNS_RECORD_KEY] == {"primary", "se1", "se2"} + + +def test_apply_skip_to_record_accumulates() -> None: + record: dict = {} + apply_skip_to_record( + record, + column_name="first", + cell_value=1, + side_effect_columns=(), + ) + apply_skip_to_record( + record, + column_name="second", + cell_value=2, + side_effect_columns=(), + ) + assert record[SKIPPED_COLUMNS_RECORD_KEY] == {"first", "second"} + assert record["first"] == 1 + assert record["second"] == 2 + + +def test_strip_skip_metadata_for_dataframe_row() -> None: + record = { + "a": 1, + SKIPPED_COLUMNS_RECORD_KEY: {"x"}, + "b": 2, + } + stripped = strip_skip_metadata_for_dataframe_row(record) + assert stripped == {"a": 1, "b": 2} + assert SKIPPED_COLUMNS_RECORD_KEY not in stripped + + +def test_strip_skip_metadata_for_dataframe_row_no_metadata() -> None: + record = {"a": 1, "b": [10, 20]} + stripped = strip_skip_metadata_for_dataframe_row(record) + assert stripped == record + assert stripped is not record + assert stripped["b"] is record["b"] + + +@pytest.mark.parametrize( + ("rows", "expected"), + [ + pytest.param( + [{"k": 1, SKIPPED_COLUMNS_RECORD_KEY: {"c"}}, {"k": 2}], + [{"k": 1}, {"k": 2}], + id="mixed", + ), + pytest.param([], [], id="empty"), + ], +) +def test_strip_skip_metadata_from_records(rows: list[dict], expected: list[dict]) -> None: + assert strip_skip_metadata_from_records(rows) == expected + + +def test_prepare_records_for_skip_metadata_round_trip_without_metadata() -> None: + rows = [{"a": 1}, {"a": 2}] + prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(rows) + assert restore_context is None + assert prepared_rows == rows + assert prepared_rows is not rows + + +def test_prepare_records_for_skip_metadata_round_trip_injects_restore_ids() -> None: + rows = [ + {"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}}, + {"a": 2}, + {"a": 3, SKIPPED_COLUMNS_RECORD_KEY: {"col_y", "col_z"}}, + ] + prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(rows) + assert restore_context is not None + assert SKIPPED_COLUMNS_RECORD_KEY not in prepared_rows[0] + assert restore_context.restore_id_column in prepared_rows[0] + assert restore_context.skipped_columns_by_source_id == { + "0": {"col_x"}, + "2": {"col_y", "col_z"}, + } + + +def test_restore_skip_metadata_uses_restore_ids_after_reorder() -> None: + old = [ + {"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}}, + {"a": 2}, + {"a": 3, SKIPPED_COLUMNS_RECORD_KEY: {"col_z"}}, + ] + prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old) + assert restore_context is not None + restore_id_column = restore_context.restore_id_column + + new = [ + {"a": 30, restore_id_column: prepared_rows[2][restore_id_column]}, + {"a": 10, restore_id_column: prepared_rows[0][restore_id_column]}, + {"a": 20, restore_id_column: prepared_rows[1][restore_id_column]}, + ] + restore_skip_metadata(new, context=restore_context, allow_resize=False) + + assert new[0][SKIPPED_COLUMNS_RECORD_KEY] == {"col_z"} + assert new[1][SKIPPED_COLUMNS_RECORD_KEY] == {"col_x"} + assert SKIPPED_COLUMNS_RECORD_KEY not in new[2] + + +def test_restore_skip_metadata_allow_resize_handles_filtered_rows() -> None: + old = [{"a": 1}, {"a": 2}] + prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old) + assert restore_context is None + + old = [ + {"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}}, + {"a": 2}, + ] + prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old) + assert restore_context is not None + restore_id_column = restore_context.restore_id_column + + new = [{"a": 20, restore_id_column: prepared_rows[1][restore_id_column]}] + restore_skip_metadata(new, context=restore_context, allow_resize=True) + + assert SKIPPED_COLUMNS_RECORD_KEY not in new[0] + + +def test_restore_skip_metadata_rejects_missing_restore_id_column() -> None: + old = [{"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}}] + _prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old) + assert restore_context is not None + + with pytest.raises(ValueError, match="must preserve the internal column"): + restore_skip_metadata([{"a": 10}], context=restore_context, allow_resize=False) diff --git a/packages/data-designer-engine/tests/engine/test_validation.py b/packages/data-designer-engine/tests/engine/test_validation.py index 9af2b6455..9bbbb13df 100644 --- a/packages/data-designer-engine/tests/engine/test_validation.py +++ b/packages/data-designer-engine/tests/engine/test_validation.py @@ -5,6 +5,7 @@ import pytest +from data_designer.config.base import SkipConfig from data_designer.config.column_configs import ( ExpressionColumnConfig, LLMCodeColumnConfig, @@ -33,6 +34,7 @@ validate_expression_references, validate_prompt_templates, validate_schema_transform_processor, + validate_skip_references, ) STUB_MODEL_ALIAS = "stub-alias" @@ -118,17 +120,19 @@ @patch("data_designer.engine.validation.validate_prompt_templates") @patch("data_designer.engine.validation.validate_code_validation") @patch("data_designer.engine.validation.validate_expression_references") +@patch("data_designer.engine.validation.validate_skip_references") @patch("data_designer.engine.validation.validate_columns_not_all_dropped") @patch("data_designer.engine.validation.validate_drop_columns_processor") @patch("data_designer.engine.validation.validate_schema_transform_processor") def test_validate_data_designer_config( - mock_validate_columns_not_all_dropped, - mock_validate_expression_references, - mock_validate_code_validation, - mock_validate_prompt_templates, - mock_validate_drop_columns_processor, - mock_validate_schema_transform_processor, -): + mock_validate_schema_transform_processor: Mock, + mock_validate_drop_columns_processor: Mock, + mock_validate_columns_not_all_dropped: Mock, + mock_validate_skip_references: Mock, + mock_validate_expression_references: Mock, + mock_validate_code_validation: Mock, + mock_validate_prompt_templates: Mock, +) -> None: mock_validate_columns_not_all_dropped.return_value = [ Violation( column="test_column", @@ -177,11 +181,20 @@ def test_validate_data_designer_config( level=ViolationLevel.ERROR, ) ] + mock_validate_skip_references.return_value = [ + Violation( + column="test_column", + type=ViolationType.SKIP_REFERENCE_MISSING, + message="test error message", + level=ViolationLevel.ERROR, + ) + ] violations = validate_data_designer_config(COLUMNS, PROCESSOR_CONFIGS, ALLOWED_REFERENCE) - assert len(violations) == 6 + assert len(violations) == 7 mock_validate_columns_not_all_dropped.assert_called_once() mock_validate_expression_references.assert_called_once() + mock_validate_skip_references.assert_called_once() mock_validate_code_validation.assert_called_once() mock_validate_prompt_templates.assert_called_once() mock_validate_drop_columns_processor.assert_called_once() @@ -349,3 +362,65 @@ def test_rich_print_violations(mock_console_print): ] ) mock_console_print.assert_called_once() + + +def test_validate_skip_references_missing_column() -> None: + columns = [ + LLMTextColumnConfig( + name="with_skip", + prompt="test {{ real_col }}", + model_alias=STUB_MODEL_ALIAS, + skip=SkipConfig(when="{{ ghost }}"), + ), + ] + violations = validate_skip_references(columns, allowed_references=["real_col"]) + assert len(violations) == 1 + assert violations[0].type == ViolationType.SKIP_REFERENCE_MISSING + assert violations[0].column == "with_skip" + + +def test_validate_skip_references_valid() -> None: + columns = [ + LLMTextColumnConfig( + name="with_skip", + prompt="test {{ gate }}", + model_alias=STUB_MODEL_ALIAS, + skip=SkipConfig(when="{{ gate == 0 }}"), + ), + ] + violations = validate_skip_references(columns, allowed_references=["gate", "with_skip"]) + assert len(violations) == 0 + + +def test_validate_skip_on_sampler_seed() -> None: + col = SamplerColumnConfig.model_construct( + name="sampler_with_skip", + column_type="sampler", + sampler_type="uniform", + params={"low": 0, "high": 10}, + skip=SkipConfig(when="{{ y }}"), + drop=False, + allow_resize=False, + propagate_skip=True, + ) + violations = validate_skip_references([col], allowed_references=["y"]) + assert len(violations) == 1 + assert violations[0].type == ViolationType.SKIP_ON_SAMPLER_SEED + assert violations[0].column == "sampler_with_skip" + + +def test_validate_skip_with_allow_resize() -> None: + col = LLMTextColumnConfig.model_construct( + name="with_skip", + column_type="llm-text", + prompt="test {{ gate }}", + model_alias=STUB_MODEL_ALIAS, + skip=SkipConfig(when="{{ gate == 0 }}"), + allow_resize=True, + drop=False, + propagate_skip=True, + ) + violations = validate_skip_references([col], allowed_references=["gate"]) + assert len(violations) == 1 + assert violations[0].type == ViolationType.SKIP_WITH_ALLOW_RESIZE + assert violations[0].column == "with_skip"