data-designer-engine 0.9.7__tar.gz → 0.10.0__tar.gz

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/PKG-INFO

@@ -1,22 +1,21 @@
 Metadata-Version: 2.5
 Name: data-designer-engine
-Version: 0.9.7
+Version: 0.10.0
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Requires-Dist: anyascii<1,>=0.3.3
 Requires-Dist: asciichartpy<2,>=1.5.25
-Requires-Dist: data-designer-config==0.9.7
+Requires-Dist: data-designer-config==0.10.0
 Requires-Dist: duckdb<2,>=1.5.0
 Requires-Dist: faker<21,>=20.1.0
 Requires-Dist: fsspec<2027,>=2026.6.0
@@ -24,7 +23,7 @@ Requires-Dist: httpx-retries<1,>=0.4.2
 Requires-Dist: httpx<1,>=0.27.2
 Requires-Dist: huggingface-hub<2,>=1.0.1
 Requires-Dist: jinja2<4,>=3.1.6
-Requires-Dist: json-repair<1,>=0.48.0
+Requires-Dist: json-repair<1,>=0.60.1
 Requires-Dist: jsonpath-rust-bindings<2,>=1.0
 Requires-Dist: jsonschema<5,>=4.0.0
 Requires-Dist: lxml<7,>=6.1.1

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/pyproject.toml

@@ -3,7 +3,7 @@ name = "data-designer-engine"
 dynamic = ["version", "dependencies"]
 description = "Generation engine for DataDesigner synthetic data generation"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 license = "Apache-2.0"
 
 classifiers = [
@@ -12,7 +12,6 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "License :: OSI Approved :: Apache Software License",
-    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
@@ -43,7 +42,7 @@ dependencies = [
     "httpx-retries>=0.4.2,<1",
     "huggingface-hub>=1.0.1,<2",
     "jinja2>=3.1.6,<4",
-    "json-repair>=0.48.0,<1",
+    "json-repair>=0.60.1,<1",
     "jsonpath-rust-bindings>=1.0,<2",
     "jsonschema>=4.0.0,<5",
     "lxml>=6.1.1,<7", # 6.1.1 fixes CVE-2026-49825

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/analysis/column_profilers/base.py

@@ -5,10 +5,10 @@ from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
+from typing import Self
 
 import pandas as pd
 from pydantic import BaseModel, model_validator
-from typing_extensions import Self
 
 from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/analysis/column_statistics.py

@@ -4,10 +4,9 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Any, TypeAlias
+from typing import TYPE_CHECKING, Any, Self, TypeAlias
 
 from pydantic import BaseModel
-from typing_extensions import Self
 
 from data_designer.config.analysis.column_statistics import (
     DEFAULT_COLUMN_STATISTICS_MAP,

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/generators/base.py

@@ -9,8 +9,9 @@ import functools
 import hashlib
 import logging
 from abc import ABC, abstractmethod
+from collections.abc import Coroutine
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Coroutine, TypeVar, overload
+from typing import TYPE_CHECKING, Any, TypeVar, overload
 
 from data_designer.config.column_configs import GenerationStrategy
 from data_designer.config.scheduling import SchedulingMetadata, SchedulingMetadataError
@@ -68,7 +69,7 @@ def _run_coroutine_sync(coro: Coroutine[Any, Any, _T]) -> _T:
     timed_out = False
     try:
         result = future.result(timeout=SYNC_BRIDGE_TIMEOUT)
-    except concurrent.futures.TimeoutError as exc:
+    except TimeoutError as exc:
         timed_out = True
         logger.warning(f"⚠️ Sync bridge timed out after {SYNC_BRIDGE_TIMEOUT}s; background thread still running")
         raise TimeoutError(f"_run_coroutine_sync timed out after {SYNC_BRIDGE_TIMEOUT}s") from exc

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/generators/custom.py

@@ -128,7 +128,7 @@ class _AsyncBridgedModelFacade:
                 return future.result(timeout=max(0.0, min(_BRIDGE_CANCEL_POLL_S, deadline - time.monotonic())))
             except concurrent.futures.CancelledError as exc:
                 raise asyncio.CancelledError from exc
-            except concurrent.futures.TimeoutError as exc:
+            except TimeoutError as exc:
                 if is_run_cancellation_requested():
                     future.cancel()
                     raise asyncio.CancelledError from exc

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/generators/embedding.py

@@ -28,8 +28,19 @@ class EmbeddingCellGenerator(ColumnGeneratorWithModel[EmbeddingColumnConfig]):
         return GenerationStrategy.CELL_BY_CELL
 
     def _prepare_embedding_inputs(self, data: dict) -> list[str]:
-        deserialized_record = deserialize_json_values(data)
-        return parse_list_string(deserialized_record[self.config.target_column])
+        cell_value = data[self.config.target_column]
+        deserialized_value = deserialize_json_values(data)[self.config.target_column]
+        if isinstance(deserialized_value, list):
+            return [str(text).strip() for text in deserialized_value]
+        if isinstance(deserialized_value, str):
+            return parse_list_string(deserialized_value)
+        if not isinstance(cell_value, str):
+            raise ValueError(
+                f"Embedding target column {self.config.target_column!r} must contain text or a list of texts, "
+                f"got {type(cell_value).__name__}"
+            )
+        # A JSON scalar such as "123" or "true" is embedded as the user wrote it.
+        return [cell_value.strip()]
 
     def generate(self, data: dict) -> dict:
         input_texts = self._prepare_embedding_inputs(data)

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/generators/expression.py

@@ -120,7 +120,16 @@ class ExpressionColumnGenerator(WithJinja2UserTemplateRendering, ColumnGenerator
         elif self.config.dtype == "float":
             return float(value)
         elif self.config.dtype == "int":
-            return int(float(value))
+            # Parse integer strings directly so values beyond 2**53 keep their exact value.
+            # Float-looking strings such as "3.0" or "1e3" still go through float().
+            try:
+                result = int(value)
+            except ValueError:
+                return int(float(value))
+            # pandas cannot build a column from an int beyond the float range, so reject it
+            # here and drop only this row, as int(float(value)) did.
+            float(result)
+            return result
         elif self.config.dtype == "bool":
             try:
                 return bool(int(float(value)))

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/generators/samplers.py

@@ -5,8 +5,9 @@ from __future__ import annotations
 
 import logging
 import random
+from collections.abc import Callable
 from functools import partial
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
 from data_designer.engine.column_generators.generators.base import FromScratchColumnGenerator, GenerationStrategy

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/generators/validation.py

@@ -94,12 +94,13 @@ class ValidationColumnGenerator(ColumnGeneratorFullColumn[ValidationColumnConfig
                 if outputs_as_dicts is None:
                     outputs_as_dicts = [{cols[0]: output.model_dump(mode="json")} for output in concatenated_outputs]
                 else:
-                    for dict_output in outputs_as_dicts:
-                        dict_output[cols[0]] = concatenated_outputs[0].model_dump(mode="json")
+                    for dict_output, output in zip(outputs_as_dicts, concatenated_outputs, strict=True):
+                        dict_output[cols[0]] = output.model_dump(mode="json")
             else:
                 outputs_as_dicts = [output.model_dump(mode="json") for output in concatenated_outputs]
 
-        validation_results = lazy.pd.DataFrame({self.config.name: outputs_as_dicts})
+        # Match the input index so results stay aligned when upstream rows were dropped.
+        validation_results = lazy.pd.DataFrame({self.config.name: outputs_as_dicts}, index=data.index)
         return lazy.pd.concat([data, validation_results], axis=1)
 
     def _validate_in_parallel(self, validator: BaseValidator, batched_records: list[list[dict]]) -> pd.DataFrame:
@@ -134,7 +135,15 @@ class ValidationColumnGenerator(ColumnGeneratorFullColumn[ValidationColumnConfig
 
     def _validate_batch(self, validator: BaseValidator, batch: list[dict]) -> ValidationResult:
         try:
-            return validator.run_validation(batch)
+            result = validator.run_validation(batch)
+            # Check each batch before results are flattened: a short batch and an excess batch could
+            # otherwise cancel out in the total and attach results to the wrong records.
+            if len(result.data) != len(batch):
+                raise DataDesignerRuntimeError(
+                    f"Validator for column {self.config.name!r} returned {len(result.data)} results "
+                    f"for a batch of {len(batch)} records."
+                )
+            return result
         except Exception as e:
             error_to_display = str(e).replace("\n", "\n  ")  # add spaces to improve readability
             logger.error(f"Batch could not be validated:\n  {error_to_display}")

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -12,6 +12,7 @@ from data_designer.config.models import ModelConfig
 from data_designer.config.run_config import JinjaRenderingEngine
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.misc import extract_keywords_from_jinja2_template
+from data_designer.config.utils.storage_schema import compile_json_storage_schema
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.engine.column_generators.utils.errors import PromptTemplateRenderError
 from data_designer.engine.column_generators.utils.judge_score_factory import (
@@ -97,7 +98,7 @@ def create_response_recipe(
         )
     if column_config.column_type == DataDesignerColumnType.LLM_STRUCTURED:
         return StructuredResponseRecipe(
-            json_schema=column_config.output_format,
+            json_schema=compile_json_storage_schema(column_config.output_format)[0],
         )
     if column_config.column_type == DataDesignerColumnType.LLM_JUDGE:
         return PydanticResponseRecipe(

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/compiler.py

@@ -9,6 +9,8 @@ from data_designer.config.column_configs import SamplerColumnConfig, SeedDataset
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.errors import InvalidConfigError
 from data_designer.config.sampler_params import UUIDSamplerParams
+from data_designer.engine.processing.processors.base import Processor
+from data_designer.engine.processing.processors.registry import ProcessorRegistry, create_default_processor_registry
 from data_designer.engine.resources.resource_provider import ResourceProvider
 from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderConfigError
 from data_designer.engine.validation import (
@@ -21,13 +23,75 @@ from data_designer.engine.validation import (
 logger = logging.getLogger(__name__)
 
 
-def compile_data_designer_config(config: DataDesignerConfig, resource_provider: ResourceProvider) -> DataDesignerConfig:
+def compile_data_designer_config(
+    config: DataDesignerConfig,
+    resource_provider: ResourceProvider,
+    processor_registry: ProcessorRegistry | None = None,
+) -> DataDesignerConfig:
+    """Resolve seed and pre-batch schemas, then validate downstream generation.
+
+    Args:
+        config: Configuration to enrich in place.
+        resource_provider: Resources used to resolve the raw seed schema.
+        processor_registry: Runtime processor implementations used to identify
+            pre-batch callbacks. Defaults to the built-in and plugin registry.
+
+    Returns:
+        The enriched, validated configuration.
+
+    Raises:
+        InvalidConfigError: If seed resolution, processor declarations, or
+            downstream validation fail.
+    """
     _resolve_and_add_seed_columns(config, resource_provider.seed_reader)
+    _apply_processor_column_modifications(config, processor_registry or create_default_processor_registry())
+    _mark_processor_dropped_seed_columns(
+        config, [column for column in config.columns if isinstance(column, SeedDatasetColumnConfig)]
+    )
     _add_internal_row_id_column_if_needed(config)
-    _validate(config)
+    _validate(config, preserve_dropped_columns=resource_provider.run_config.preserve_dropped_columns)
     return config
 
 
+def _apply_processor_column_modifications(config: DataDesignerConfig, registry: ProcessorRegistry) -> None:
+    """Adjusts columns according to columns_added and columns_removed declared by processors."""
+    for processor in config.processors or []:
+        if not processor.columns_added and not processor.columns_removed:
+            continue
+        processor_cls = registry.get_for_config_type(type(processor))
+        if processor_cls.process_before_batch is Processor.process_before_batch:
+            continue
+        if processor.columns_removed:
+            current_columns = {col.name for col in config.columns}
+            for col_name in processor.columns_removed:
+                if col_name not in current_columns:
+                    raise InvalidConfigError(
+                        f"🛑 Processor '{processor.name}' cannot remove column '{col_name}' because it does not exist."
+                    )
+            for column in config.columns:
+                if column.name in processor.columns_removed and not isinstance(column, SeedDatasetColumnConfig):
+                    raise InvalidConfigError(
+                        f"🛑 Pre-batch processor '{processor.name}' can only remove seed columns; "
+                        f"'{column.name}' is a generated column."
+                    )
+            removed_set = set(processor.columns_removed)
+            config.columns = [col for col in config.columns if col.name not in removed_set]
+
+        if processor.columns_added:
+            if config.seed_config is None:
+                raise InvalidConfigError(
+                    f"🛑 Processor '{processor.name}' specifies 'columns_added', but no seed dataset is configured."
+                )
+            existing_columns = set(_get_allowed_references(config))
+            for col_name in processor.columns_added:
+                if col_name in existing_columns:
+                    raise InvalidConfigError(
+                        f"🛑 Processor '{processor.name}' adds column '{col_name}' which collides with an existing column."
+                    )
+                config.columns.append(SeedDatasetColumnConfig(name=col_name))
+                existing_columns.add(col_name)
+
+
 def _resolve_and_add_seed_columns(config: DataDesignerConfig, seed_reader: SeedReader | None) -> None:
     """Fetches the seed dataset column names, ensures there are no conflicts
     with other columns, and adds seed column configs to the DataDesignerConfig.
@@ -50,7 +114,6 @@ def _resolve_and_add_seed_columns(config: DataDesignerConfig, seed_reader: SeedR
 
     seed_columns = [SeedDatasetColumnConfig(name=col_name) for col_name in seed_col_names]
     config.columns.extend(seed_columns)
-    _mark_processor_dropped_seed_columns(config, seed_columns)
 
 
 def _mark_processor_dropped_seed_columns(
@@ -58,8 +121,8 @@ def _mark_processor_dropped_seed_columns(
 ) -> None:
     """Marks seed columns targeted by a drop processor with `drop=True`.
 
-    `DataDesignerConfigBuilder.add_processor` does this for configured columns, but seed
-    columns don't exist yet at that point — they're resolved here. Without this, a
+    `DataDesignerConfigBuilder.build` does this for configured columns, but seed columns
+    aren't known to the builder — they're resolved here. Without this, a
     processor-dropped seed column is still handed to the dataset profiler, which either
     reports a column the emitted dataset no longer has or fails outright when
     `preserve_dropped_columns=False`.
@@ -84,9 +147,9 @@ def _add_internal_row_id_column_if_needed(config: DataDesignerConfig) -> None:
         config: The DataDesigner configuration to potentially modify.
     """
     has_sampler_column = any(isinstance(col, SamplerColumnConfig) for col in config.columns)
-    has_seed_dataset_column = any(isinstance(col, SeedDatasetColumnConfig) for col in config.columns)
+    has_seed_dataset = config.seed_config is not None
 
-    if not has_sampler_column and not has_seed_dataset_column:
+    if not has_sampler_column and not has_seed_dataset:
         logger.warning(
             "🔔 No sampler column or seed dataset detected. Adding UUID column '_internal_row_id' (marked for drop) to enable generation."
         )
@@ -99,12 +162,13 @@ def _add_internal_row_id_column_if_needed(config: DataDesignerConfig) -> None:
         config.columns.insert(0, id_column)
 
 
-def _validate(config: DataDesignerConfig) -> None:
+def _validate(config: DataDesignerConfig, *, preserve_dropped_columns: bool = True) -> None:
     allowed_references = _get_allowed_references(config)
     violations = validate_data_designer_config(
         columns=config.columns,
         processor_configs=config.processors or [],
         allowed_references=allowed_references,
+        preserve_dropped_columns=preserve_dropped_columns,
     )
     rich_print_violations(violations)
     if len([v for v in violations if v.level == ViolationLevel.ERROR]) > 0:

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/dataset_builders/async_scheduler.py

@@ -10,10 +10,10 @@ import logging
 import time
 import uuid
 from collections import Counter, defaultdict, deque
-from collections.abc import Coroutine, Mapping
+from collections.abc import Callable, Coroutine, Mapping
 from dataclasses import dataclass
 from threading import Event
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 import data_designer.lazy_heavy_imports as lazy
 from data_designer.config.column_configs import ExpressionColumnConfig, GenerationStrategy
@@ -59,10 +59,7 @@ from data_designer.engine.dataset_builders.scheduling.task_admission import (
 from data_designer.engine.dataset_builders.scheduling.task_model import SliceRef, Task, TaskTrace
 from data_designer.engine.dataset_builders.scheduling.task_policies import BoundedBorrowTaskAdmissionPolicyConfig
 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.skip_tracker import apply_skip_to_record
 from data_designer.engine.errors import DataDesignerError
 from data_designer.engine.models.clients.errors import ProviderError
 from data_designer.engine.models.errors import (
@@ -1560,7 +1557,7 @@ class AsyncTaskScheduler:
     async def _wait_before_retryable_resalvage(self) -> None:
         """Pace repeated preserved retryable salvage cycles to avoid a hot loop."""
         self._wake_event.clear()
-        with contextlib.suppress(asyncio.TimeoutError):
+        with contextlib.suppress(TimeoutError):
             await asyncio.wait_for(self._wake_event.wait(), timeout=self._retryable_resalvage_delay_seconds())
         self._raise_if_fatal_worker_error()
 
@@ -2246,8 +2243,7 @@ class AsyncTaskScheduler:
             # empty DataFrame which would yield zero values and fail
             # ``update_batch``.
             if self._buffer_manager is not None:
-                records = [self._buffer_manager.get_row(task.row_group, ri) for ri in range(rg_size)]
-                input_df = lazy.pd.DataFrame(records)
+                input_df = self._buffer_manager.get_dataframe(task.row_group, row_indices=list(range(rg_size)))
             else:
                 input_df = lazy.pd.DataFrame(index=range(rg_size))
             result_df = await self._run_generator_call(
@@ -2292,8 +2288,7 @@ class AsyncTaskScheduler:
         if self._buffer_manager is not None:
             pre_dropped: set[int] = {ri for ri in range(rg_size) if self._buffer_manager.is_dropped(task.row_group, ri)}
             active_row_indices = [ri for ri in range(rg_size) if ri not in pre_dropped]
-            records = [self._buffer_manager.get_row(task.row_group, ri) for ri in active_row_indices]
-            input_df = lazy.pd.DataFrame(records, index=active_row_indices) if records else lazy.pd.DataFrame()
+            input_df = self._buffer_manager.get_dataframe(task.row_group, row_indices=active_row_indices)
         else:
             pre_dropped = set()
             active_row_indices = list(range(rg_size))
@@ -2413,7 +2408,6 @@ class AsyncTaskScheduler:
 
         if self._buffer_manager is not None:
             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] = []
             active_row_indices: list[int] = []
 
             # Skip evaluation only applies to single-column configs.
@@ -2431,14 +2425,9 @@ class AsyncTaskScheduler:
                     pre_skipped.add(ri)
                     continue
 
-                active_rows_data.append(record)
                 active_row_indices.append(ri)
 
-            batch_df = (
-                lazy.pd.DataFrame(strip_skip_metadata_from_records(active_rows_data), index=active_row_indices)
-                if active_rows_data
-                else lazy.pd.DataFrame()
-            )
+            batch_df = self._buffer_manager.get_dataframe(task.row_group, row_indices=active_row_indices)
         else:
             batch_df = lazy.pd.DataFrame()
             pre_dropped = set()

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/dataset_builders/dataset_builder.py

@@ -12,9 +12,10 @@ import logging
 import os
 import time
 import uuid
+from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from pydantic import ValidationError
 
@@ -25,8 +26,8 @@ from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
     ProcessorConfig,
-    ProcessorType,
 )
+from data_designer.config.utils.storage_schema import JSON_STORAGE_RUN_KEY, JsonStoragePlan, compile_json_storage_schema
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.config.version import get_library_version
 from data_designer.engine.column_generators.generators.base import (
@@ -67,6 +68,7 @@ from data_designer.engine.storage.artifact_storage import (
     ResumeMode,
 )
 from data_designer.engine.storage.media_storage import StorageMode
+from data_designer.engine.validation import resolve_processor_dropped_columns
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -194,14 +196,39 @@ class DatasetBuilder:
         # when a run produces 0 records due to deterministic failures.
         self._first_non_retryable_error: Exception | 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)
+        self._data_designer_config = compile_data_designer_config(
+            data_designer_config, resource_provider, self._registry.processors
+        )
+        self._column_configs = compile_dataset_builder_column_configs(
+            self._data_designer_config,
+            seed_column_names=(
+                resource_provider.seed_reader.get_column_names() if resource_provider.seed_reader is not None else None
+            ),
+        )
         processors = self._initialize_processors(self._data_designer_config.processors or [])
         self._processor_runner = ProcessorRunner(
             processors=processors,
             artifact_storage=resource_provider.artifact_storage,
         )
         self._validate_column_configs()
+        columns = {}
+        if resource_provider.seed_reader is not None:
+            seed_plan = resource_provider.seed_reader.get_storage_plan()
+            if isinstance(seed_plan, JsonStoragePlan):
+                seed_column_names = [
+                    column.name
+                    for column in self.single_column_configs
+                    if column.column_type == DataDesignerColumnType.SEED_DATASET
+                ]
+                columns.update(seed_plan.restrict_to_columns(seed_column_names).columns)
+        for column in self.single_column_configs:
+            if column.storage_encoding == "json":
+                columns[column.name] = [[]]
+            elif getattr(column, "column_type", None) == "llm-structured":
+                _, paths = compile_json_storage_schema(column.output_format)
+                if paths:
+                    columns[column.name] = paths
+        self.artifact_storage.json_storage = JsonStoragePlan(columns=columns)
 
     @property
     def artifact_storage(self) -> ArtifactStorage:
@@ -415,6 +442,8 @@ class DatasetBuilder:
             return None
 
         post_generation_state = metadata.get("post_generation_state")
+        if not self._json_storage_policy_matches(metadata):
+            raise DatasetGenerationError("Cannot resume: JSON storage contract differs from the saved run")
         if post_generation_state == "started":
             raise DatasetGenerationError(
                 "🛑 Cannot resume: process_after_generation started but did not complete for this dataset. "
@@ -485,6 +514,9 @@ class DatasetBuilder:
             allow_holes=True,
         )
 
+        if not self._json_storage_policy_matches(metadata):
+            raise DatasetGenerationError("Cannot resume: JSON storage contract differs from the saved run")
+
         if num_records < actual_num_records:
             raise DatasetGenerationError(
                 f"🛑 Cannot resume: num_records={num_records} is less than the {actual_num_records} "
@@ -678,6 +710,8 @@ class DatasetBuilder:
 
             if not self._dropped_column_artifact_policy_matches(metadata):
                 return _ConfigCompatibility.INCOMPATIBLE
+            if not self._json_storage_policy_matches(metadata):
+                return _ConfigCompatibility.INCOMPATIBLE
 
             stored_hash = metadata.get("config_hash")
             stored_version = metadata.get("config_hash_version")
@@ -721,6 +755,13 @@ class DatasetBuilder:
             )
             return _ConfigCompatibility.INCOMPATIBLE
 
+    def _json_storage_policy_matches(self, metadata: dict[str, Any]) -> bool:
+        try:
+            stored = JsonStoragePlan.model_validate(metadata.get(JSON_STORAGE_RUN_KEY, {}))
+        except ValueError as exc:
+            raise DatasetGenerationError("Invalid JSON storage contract in run metadata") from exc
+        return stored == self.artifact_storage.json_storage
+
     def _dropped_column_artifact_policy_matches(self, metadata: dict[str, Any]) -> bool:
         """Return whether stored dropped-column artifact behavior matches this run.
 
@@ -1014,7 +1055,13 @@ class DatasetBuilder:
 
     def process_preview(self, dataset: pd.DataFrame) -> pd.DataFrame:
         df = self._processor_runner.run_post_batch(dataset.copy(), current_batch_number=None)
-        return self._processor_runner.run_after_generation_on_df(df)
+        df = self._processor_runner.run_after_generation_on_df(df)
+        storage_plan = self.artifact_storage.json_storage.restrict_to_columns(list(df.columns))
+        if storage_plan.columns:
+            df.attrs[JSON_STORAGE_RUN_KEY] = storage_plan.model_dump()
+        else:
+            df.attrs.pop(JSON_STORAGE_RUN_KEY, None)
+        return df
 
     def _has_image_columns(self) -> bool:
         """Check if config has any image generation columns."""
@@ -1056,6 +1103,9 @@ class DatasetBuilder:
     def _initialize_processors(self, processor_configs: list[ProcessorConfig]) -> list[Processor]:
         # Check columns marked for drop
         columns_to_drop = [config.name for config in self.single_column_configs if config.drop]
+        # Columns an explicit "drop columns" processor already handles take precedence, including glob targets
+        explicitly_dropped = resolve_processor_dropped_columns(self.single_column_configs, processor_configs)
+        columns_to_drop = [name for name in columns_to_drop if name not in explicitly_dropped]
 
         processors: list[Processor] = []
         for config in processor_configs:
@@ -1066,12 +1116,6 @@ class DatasetBuilder:
                 )
             )
 
-            # Manually included "drop columns" processor takes precedence
-            if config.processor_type == ProcessorType.DROP_COLUMNS:
-                for column in config.column_names:
-                    if column in columns_to_drop:
-                        columns_to_drop.remove(column)
-
         # If there are still columns marked for drop, add the "drop columns" processor to drop them
         if len(columns_to_drop) > 0:
             processors.append(

{data_designer_engine-0.9.7 → data_designer_engine-0.10.0}/src/data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+from data_designer.config.column_configs import SeedDatasetColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.engine.dataset_builders.multi_column_configs import (
@@ -14,7 +15,22 @@ from data_designer.engine.dataset_builders.utils.errors import ConfigCompilation
 from data_designer.engine.dataset_builders.utils.execution_graph import topologically_sort_column_configs
 
 
-def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[DatasetBuilderColumnConfigT]:
+def compile_dataset_builder_column_configs(
+    config: DataDesignerConfig, *, seed_column_names: list[str] | None = None
+) -> list[DatasetBuilderColumnConfigT]:
+    """Group generation tasks while preserving raw inputs for pre-batch processors.
+
+    Args:
+        config: Validated configuration with the schema after pre-batch processing.
+        seed_column_names: Raw seed names from the reader, including names removed
+            by pre-batch processors. Required to retain those seed-loading tasks.
+
+    Returns:
+        Ordered seed, sampler, and generated column tasks.
+
+    Raises:
+        ConfigCompilationError: If seed columns have no seed configuration.
+    """
     seed_column_configs = []
     sampler_column_configs = []
     generated_column_configs = []
@@ -27,6 +43,12 @@ def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[D
         else:
             generated_column_configs.append(column_config)
 
+    # Raw seed columns must still be read before processors remove or replace them.
+    effective_seed_names = {column.name for column in seed_column_configs}
+    seed_column_configs.extend(
+        SeedDatasetColumnConfig(name=name) for name in seed_column_names or [] if name not in effective_seed_names
+    )
+
     compiled_column_configs = []
 
     if len(seed_column_configs) > 0: