data-designer 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

data_designer/engine/models/parsers/types.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Optional, Protocol, Type, runtime_checkable
+from typing import Any, Protocol, runtime_checkable
 
 from lxml.etree import _Element
 from pydantic import BaseModel, Field
@@ -30,7 +30,7 @@ class LLMStructuredResponse(BaseModel):
         out.parsed = out.parsed[-n:]
         return out
 
-    def filter(self, block_types: list[Type[BaseModel]]) -> Self:
+    def filter(self, block_types: list[type[BaseModel]]) -> Self:
         out = self.model_copy()
         out.parsed = [b for b in out.parsed if isinstance(b, tuple(block_types))]
         return out
@@ -44,7 +44,7 @@ class TagParser(Protocol):
     element, do some computation, and return some kind of structured
     output, represented as a subclass of Pydantic `BaseModel`.
     This protocol implementation can cover both classes as well
-    as curried fuctions as parsers (e.g. `partial`).
+    as curried functions as parsers (e.g. `partial`).
     """
 
     def __call__(self, element: _Element) -> BaseModel: ...
@@ -69,7 +69,7 @@ class TextBlock(BaseModel):
 
 class CodeBlock(BaseModel):
     code: str
-    code_lang: Optional[str] = None
+    code_lang: str | None = None
 
 
 class StructuredDataBlock(BaseModel):

data_designer/engine/models/registry.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import logging
 
-from data_designer.config.models import ModelConfig
+from data_designer.config.models import GenerationType, ModelConfig
 from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
 from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.litellm_overrides import apply_litellm_patches
@@ -73,7 +73,7 @@ class ModelRegistry:
         model_config = self.get_model_config(model_alias=model_alias)
         return self._model_provider_registry.get_provider(model_config.provider)
 
-    def run_health_check(self, model_aliases: set[str]) -> None:
+    def run_health_check(self, model_aliases: list[str]) -> None:
         logger.info("🩺 Running health checks for models...")
         for model_alias in model_aliases:
             model = self.get_model(model_alias=model_alias)
@@ -81,15 +81,24 @@ class ModelRegistry:
                 f"  |-- 👀 Checking {model.model_name!r} in provider named {model.model_provider_name!r} for model alias {model.model_alias!r}..."
             )
             try:
-                model.generate(
-                    prompt="Hello!",
-                    parser=lambda x: x,
-                    system_prompt="You are a helpful assistant.",
-                    max_correction_steps=0,
-                    max_conversation_restarts=0,
-                    skip_usage_tracking=True,
-                    purpose="running health checks",
-                )
+                if model.model_generation_type == GenerationType.EMBEDDING:
+                    model.generate_text_embeddings(
+                        input_texts=["Hello!"],
+                        skip_usage_tracking=True,
+                        purpose="running health checks",
+                    )
+                elif model.model_generation_type == GenerationType.CHAT_COMPLETION:
+                    model.generate(
+                        prompt="Hello!",
+                        parser=lambda x: x,
+                        system_prompt="You are a helpful assistant.",
+                        max_correction_steps=0,
+                        max_conversation_restarts=0,
+                        skip_usage_tracking=True,
+                        purpose="running health checks",
+                    )
+                else:
+                    raise ValueError(f"Unsupported generation type: {model.model_generation_type}")
                 logger.info("  |-- ✅ Passed!")
             except Exception as e:
                 logger.error("  |-- ❌ Failed!")

data_designer/engine/models/usage.py

@@ -11,20 +11,20 @@ logger = logging.getLogger(__name__)
 
 
 class TokenUsageStats(BaseModel):
-    prompt_tokens: int = 0
-    completion_tokens: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
 
     @computed_field
     def total_tokens(self) -> int:
-        return self.prompt_tokens + self.completion_tokens
+        return self.input_tokens + self.output_tokens
 
     @property
     def has_usage(self) -> bool:
         return self.total_tokens > 0
 
-    def extend(self, *, prompt_tokens: int, completion_tokens: int) -> None:
-        self.prompt_tokens += prompt_tokens
-        self.completion_tokens += completion_tokens
+    def extend(self, *, input_tokens: int, output_tokens: int) -> None:
+        self.input_tokens += input_tokens
+        self.output_tokens += output_tokens
 
 
 class RequestUsageStats(BaseModel):
@@ -56,9 +56,7 @@ class ModelUsageStats(BaseModel):
         self, *, token_usage: TokenUsageStats | None = None, request_usage: RequestUsageStats | None = None
     ) -> None:
         if token_usage is not None:
-            self.token_usage.extend(
-                prompt_tokens=token_usage.prompt_tokens, completion_tokens=token_usage.completion_tokens
-            )
+            self.token_usage.extend(input_tokens=token_usage.input_tokens, output_tokens=token_usage.output_tokens)
         if request_usage is not None:
             self.request_usage.extend(
                 successful_requests=request_usage.successful_requests, failed_requests=request_usage.failed_requests

data_designer/engine/processing/ginja/ast.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections import deque
-from typing import Optional, Type
 
 from jinja2 import nodes as j_nodes
 
@@ -33,7 +32,7 @@ def ast_max_depth(node: j_nodes.Node) -> int:
     return max_depth
 
 
-def ast_descendant_count(ast: j_nodes.Node, only_type: Optional[Type[j_nodes.Node]] = None) -> int:
+def ast_descendant_count(ast: j_nodes.Node, only_type: type[j_nodes.Node] | None = None) -> int:
     """Count the number of nodes which descend from the given node.
 
     Args:

data_designer/engine/processing/processors/drop_columns.py

@@ -17,7 +17,7 @@ class DropColumnsProcessor(Processor[DropColumnsProcessorConfig]):
     @staticmethod
     def metadata() -> ConfigurableTaskMetadata:
         return ConfigurableTaskMetadata(
-            name="drop_columns",
+            name="drop_columns_processor",
             description="Drop columns from the input dataset.",
             required_resources=None,
         )

data_designer/engine/processing/processors/registry.py

@@ -5,9 +5,11 @@ from data_designer.config.base import ConfigBase
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
     ProcessorType,
+    SchemaTransformProcessorConfig,
 )
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
+from data_designer.engine.processing.processors.schema_transform import SchemaTransformProcessor
 from data_designer.engine.registry.base import TaskRegistry
 
 
@@ -16,5 +18,6 @@ class ProcessorRegistry(TaskRegistry[str, Processor, ConfigBase]): ...
 
 def create_default_processor_registry() -> ProcessorRegistry:
     registry = ProcessorRegistry()
+    registry.register(ProcessorType.SCHEMA_TRANSFORM, SchemaTransformProcessor, SchemaTransformProcessorConfig, False)
     registry.register(ProcessorType.DROP_COLUMNS, DropColumnsProcessor, DropColumnsProcessorConfig, False)
     return registry

data_designer/engine/processing/processors/schema_transform.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+
+import pandas as pd
+
+from data_designer.config.processors import SchemaTransformProcessorConfig
+from data_designer.engine.configurable_task import ConfigurableTaskMetadata
+from data_designer.engine.dataset_builders.artifact_storage import BatchStage
+from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
+from data_designer.engine.processing.processors.base import Processor
+from data_designer.engine.processing.utils import deserialize_json_values
+
+logger = logging.getLogger(__name__)
+
+
+class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[SchemaTransformProcessorConfig]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="schema_transform_processor",
+            description="Generate dataset with transformed schema using a Jinja2 template.",
+            required_resources=None,
+        )
+
+    @property
+    def template_as_str(self) -> str:
+        return json.dumps(self.config.template)
+
+    def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
+        self.prepare_jinja2_template_renderer(self.template_as_str, data.columns.to_list())
+        formatted_records = [
+            json.loads(self.render_template(deserialize_json_values(record)).replace("\n", "\\n"))
+            for record in data.to_dict(orient="records")
+        ]
+        formatted_data = pd.DataFrame(formatted_records)
+        if current_batch_number is not None:
+            self.artifact_storage.write_batch_to_parquet_file(
+                batch_number=current_batch_number,
+                dataframe=formatted_data,
+                batch_stage=BatchStage.PROCESSORS_OUTPUTS,
+                subfolder=self.config.name,
+            )
+        else:
+            self.artifact_storage.write_parquet_file(
+                parquet_file_name=f"{self.config.name}.parquet",
+                dataframe=formatted_data,
+                batch_stage=BatchStage.PROCESSORS_OUTPUTS,
+            )
+
+        return data

data_designer/engine/processing/utils.py

@@ -1,9 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import ast
 import json
 import logging
-from typing import Any, TypeVar, Union, overload
+import re
+from typing import Any, TypeVar, overload
 
 import pandas as pd
 
@@ -25,7 +27,7 @@ def concat_datasets(datasets: list[pd.DataFrame]) -> pd.DataFrame:
 # Overloads to help static type checker better understand
 # the input/output types of the deserialize_json_values function.
 @overload
-def deserialize_json_values(data: str) -> Union[dict[str, Any], list[Any], Any]: ...
+def deserialize_json_values(data: str) -> dict[str, Any] | list[Any] | Any: ...
 
 
 @overload
@@ -100,6 +102,42 @@ def deserialize_json_values(data):
         return data
 
 
+def parse_list_string(text: str) -> list[str]:
+    """Parse a list from a string, handling JSON arrays, Python lists, and trailing commas."""
+    text = text.strip()
+
+    # Try JSON first
+    try:
+        list_obj = json.loads(text)
+        if isinstance(list_obj, list):
+            return _clean_whitespace(list_obj)
+    except json.JSONDecodeError:
+        pass
+
+    # Remove trailing commas before closing brackets (common in JSON-like strings)
+    text_cleaned = re.sub(r",\s*]", "]", text)
+    text_cleaned = re.sub(r",\s*}", "}", text_cleaned)
+
+    # Try JSON again with cleaned text
+    try:
+        return _clean_whitespace(json.loads(text_cleaned))
+    except json.JSONDecodeError:
+        pass
+
+    # Try Python literal eval (handles single quotes)
+    try:
+        return _clean_whitespace(ast.literal_eval(text_cleaned))
+    except (ValueError, SyntaxError):
+        pass
+
+    # If all else fails, return the original text
+    return [text.strip()]
+
+
+def _clean_whitespace(texts: list[str]) -> list[str]:
+    return [text.strip() for text in texts]
+
+
 def _verify_columns_are_unique(datasets: list[pd.DataFrame]) -> None:
     joined_columns = set()
     for df in datasets:

data_designer/engine/registry/base.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import threading
-from typing import Any, Generic, Type, TypeVar
+from typing import Any, Generic, TypeVar
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.utils.type_helpers import StrEnum
@@ -16,14 +16,14 @@ TaskConfigT = TypeVar("TaskConfigT", bound=ConfigBase)
 
 class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
     # registered type name -> type
-    _registry: dict[EnumNameT, Type[TaskT]] = {}
+    _registry: dict[EnumNameT, type[TaskT]] = {}
     # type -> registered type name
-    _reverse_registry: dict[Type[TaskT], EnumNameT] = {}
+    _reverse_registry: dict[type[TaskT], EnumNameT] = {}
 
     # registered type name -> config type
-    _config_registry: dict[EnumNameT, Type[TaskConfigT]] = {}
+    _config_registry: dict[EnumNameT, type[TaskConfigT]] = {}
     # config type -> registered type name
-    _reverse_config_registry: dict[Type[TaskConfigT], EnumNameT] = {}
+    _reverse_config_registry: dict[type[TaskConfigT], EnumNameT] = {}
 
     # all registries are singletons
     _instance = None
@@ -33,8 +33,8 @@ class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
     def register(
         cls,
         name: EnumNameT,
-        task: Type[TaskT],
-        config: Type[TaskConfigT],
+        task: type[TaskT],
+        config: type[TaskConfigT],
         raise_on_collision: bool = False,
     ) -> None:
         if cls._has_been_registered(name):
@@ -52,22 +52,22 @@ class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
             cls._reverse_config_registry[config] = name
 
     @classmethod
-    def get_task_type(cls, name: EnumNameT) -> Type[TaskT]:
+    def get_task_type(cls, name: EnumNameT) -> type[TaskT]:
         cls._raise_if_not_registered(name, cls._registry)
         return cls._registry[name]
 
     @classmethod
-    def get_config_type(cls, name: EnumNameT) -> Type[TaskConfigT]:
+    def get_config_type(cls, name: EnumNameT) -> type[TaskConfigT]:
         cls._raise_if_not_registered(name, cls._config_registry)
         return cls._config_registry[name]
 
     @classmethod
-    def get_registered_name(cls, task: Type[TaskT]) -> EnumNameT:
+    def get_registered_name(cls, task: type[TaskT]) -> EnumNameT:
         cls._raise_if_not_registered(task, cls._reverse_registry)
         return cls._reverse_registry[task]
 
     @classmethod
-    def get_for_config_type(cls, config: Type[TaskConfigT]) -> Type[TaskT]:
+    def get_for_config_type(cls, config: type[TaskConfigT]) -> type[TaskT]:
         cls._raise_if_not_registered(config, cls._reverse_config_registry)
         name = cls._reverse_config_registry[config]
         return cls.get_task_type(name)
@@ -77,7 +77,7 @@ class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
         return name in cls._registry
 
     @classmethod
-    def _raise_if_not_registered(cls, key: EnumNameT | Type[TaskT] | Type[TaskConfigT], mapping: dict) -> None:
+    def _raise_if_not_registered(cls, key: EnumNameT | type[TaskT] | type[TaskConfigT], mapping: dict) -> None:
         if not (isinstance(key, StrEnum) or isinstance(key, str)):
             cls._raise_if_not_type(key)
         if key not in mapping:

data_designer/engine/sampling_gen/constraints.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Type
 
 import numpy as np
 import pandas as pd
@@ -91,5 +90,5 @@ CONSTRAINT_TYPE_TO_CHECKER = {
 }
 
 
-def get_constraint_checker(constraint_type: ConstraintType) -> Type[ConstraintChecker]:
+def get_constraint_checker(constraint_type: ConstraintType) -> type[ConstraintChecker]:
     return CONSTRAINT_TYPE_TO_CHECKER[ConstraintType(constraint_type)]

data_designer/engine/sampling_gen/data_sources/base.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Optional, Type, TypeVar, Union
+from typing import Any, Generic, TypeVar
 
 import numpy as np
 import pandas as pd
@@ -45,7 +45,7 @@ class PassthroughMixin:
         return series
 
     @staticmethod
-    def validate_data_conversion(convert_to: Optional[str]) -> None:
+    def validate_data_conversion(convert_to: str | None) -> None:
         pass
 
 
@@ -71,7 +71,7 @@ class TypeConversionMixin:
         return series
 
     @staticmethod
-    def postproc(series: pd.Series, convert_to: Optional[str]) -> pd.Series:
+    def postproc(series: pd.Series, convert_to: str | None) -> pd.Series:
         if convert_to is not None:
             if convert_to == "int":
                 series = series.round()
@@ -79,18 +79,18 @@ class TypeConversionMixin:
         return series
 
     @staticmethod
-    def validate_data_conversion(convert_to: Optional[str]) -> None:
+    def validate_data_conversion(convert_to: str | None) -> None:
         if convert_to is not None and convert_to not in ["float", "int", "str"]:
             raise ValueError(f"Invalid `convert_to` value: {convert_to}. Must be one of: [float, int, str]")
 
 
 class DatetimeFormatMixin:
     @staticmethod
-    def preproc(series: pd.Series, convert_to: Optional[str]) -> pd.Series:
+    def preproc(series: pd.Series, convert_to: str | None) -> pd.Series:
         return series
 
     @staticmethod
-    def postproc(series: pd.Series, convert_to: Optional[str]) -> pd.Series:
+    def postproc(series: pd.Series, convert_to: str | None) -> pd.Series:
         if convert_to is not None:
             return series.dt.strftime(convert_to)
         if series.dt.month.nunique() == 1:
@@ -104,7 +104,7 @@ class DatetimeFormatMixin:
         return series.apply(lambda dt: dt.isoformat()).astype(str)
 
     @staticmethod
-    def validate_data_conversion(convert_to: Optional[str]) -> None:
+    def validate_data_conversion(convert_to: str | None) -> None:
         if convert_to is not None:
             try:
                 pd.to_datetime(pd.to_datetime("2012-12-21").strftime(convert_to))
@@ -121,7 +121,7 @@ class DataSource(ABC, Generic[GenericParamsT]):
     def __init__(
         self,
         params: GenericParamsT,
-        random_state: Optional[RadomStateT] = None,
+        random_state: RadomStateT | None = None,
         **kwargs,
     ):
         self.rng = check_random_state(random_state)
@@ -130,7 +130,7 @@ class DataSource(ABC, Generic[GenericParamsT]):
         self._validate()
 
     @classmethod
-    def get_param_type(cls) -> Type[GenericParamsT]:
+    def get_param_type(cls) -> type[GenericParamsT]:
         return cls.__orig_bases__[-1].__args__[0]
 
     @abstractmethod
@@ -138,7 +138,7 @@ class DataSource(ABC, Generic[GenericParamsT]):
         self,
         dataframe: pd.DataFrame,
         column_name: str,
-        index: Optional[list[int]] = None,
+        index: list[int] | None = None,
     ) -> pd.DataFrame: ...
 
     @staticmethod
@@ -147,11 +147,11 @@ class DataSource(ABC, Generic[GenericParamsT]):
 
     @staticmethod
     @abstractmethod
-    def postproc(series: pd.Series, convert_to: Optional[str]) -> pd.Series: ...
+    def postproc(series: pd.Series, convert_to: str | None) -> pd.Series: ...
 
     @staticmethod
     @abstractmethod
-    def validate_data_conversion(convert_to: Optional[str]) -> None: ...
+    def validate_data_conversion(convert_to: str | None) -> None: ...
 
     def get_required_column_names(self) -> tuple[str, ...]:
         return tuple()
@@ -182,7 +182,7 @@ class Sampler(DataSource[GenericParamsT], ABC):
         self,
         dataframe: pd.DataFrame,
         column_name: str,
-        index: Optional[list[int]] = None,
+        index: list[int] | None = None,
     ) -> pd.DataFrame:
         index = slice(None) if index is None else index
 
@@ -208,7 +208,7 @@ class Sampler(DataSource[GenericParamsT], ABC):
 class ScipyStatsSampler(Sampler[GenericParamsT], ABC):
     @property
     @abstractmethod
-    def distribution(self) -> Union[stats.rv_continuous, stats.rv_discrete]: ...
+    def distribution(self) -> stats.rv_continuous | stats.rv_discrete: ...
 
     def sample(self, num_samples: int) -> NumpyArray1dT:
         return self.distribution.rvs(size=num_samples, random_state=self.rng)

data_designer/engine/sampling_gen/entities/phone_number.py

@@ -3,7 +3,6 @@
 
 import random
 from pathlib import Path
-from typing import Optional
 
 import pandas as pd
 from pydantic import BaseModel, Field, field_validator
@@ -13,7 +12,7 @@ ZIPCODE_AREA_CODE_MAP = dict(zip(ZIP_AREA_CODE_DATA["zipcode"], ZIP_AREA_CODE_DA
 ZIPCODE_POPULATION_MAP = dict(zip(ZIP_AREA_CODE_DATA["zipcode"], ZIP_AREA_CODE_DATA["count"]))
 
 
-def get_area_code(zip_prefix: Optional[str] = None) -> str:
+def get_area_code(zip_prefix: str | None = None) -> str:
     """
     Sample an area code for the given ZIP code prefix, population-weighted.
 

data_designer/engine/sampling_gen/people_gen.py

@@ -8,12 +8,12 @@ import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, TypeAlias
 
 import pandas as pd
 from faker import Faker
 
-from data_designer.config.utils.constants import AVAILABLE_LOCALES, DEFAULT_AGE_RANGE
+from data_designer.config.utils.constants import DEFAULT_AGE_RANGE
 from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
 from data_designer.engine.sampling_gen.entities.dataset_based_person_fields import PERSONA_FIELDS, PII_FIELDS
 from data_designer.engine.sampling_gen.entities.person import (
@@ -27,17 +27,13 @@ if TYPE_CHECKING:
     from data_designer.engine.sampling_gen.schema import DataSchema
 
 
-EngineT = Union[Faker, ManagedDatasetGenerator]
+EngineT: TypeAlias = Faker | ManagedDatasetGenerator
 
 
 class PeopleGen(ABC):
     """Unified interface for generating people data."""
 
     def __init__(self, engine: EngineT, locale: str):
-        if locale not in AVAILABLE_LOCALES:
-            raise ValueError(
-                f"Locale {locale} is not a supported locale.Supported locales: {', '.join(AVAILABLE_LOCALES)}"
-            )
         self.locale = locale
         self._engine = engine
 

data_designer/engine/validators/base.py

@@ -2,14 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Iterator, Optional
+from typing import Iterator
 
 from pydantic import BaseModel, ConfigDict
 from typing_extensions import Self
 
 
 class ValidationOutput(BaseModel):
-    is_valid: Optional[bool]
+    is_valid: bool | None
     model_config = ConfigDict(extra="allow")
 
 

data_designer/interface/data_designer.py

@@ -249,6 +249,17 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         except Exception as e:
             raise DataDesignerProfilingError(f"🛑 Error profiling preview dataset: {e}")
 
+        if builder.artifact_storage.processors_outputs_path.exists():
+            processor_artifacts = {
+                processor_config.name: pd.read_parquet(
+                    builder.artifact_storage.processors_outputs_path / f"{processor_config.name}.parquet",
+                    dtype_backend="pyarrow",
+                ).to_dict(orient="records")
+                for processor_config in config_builder.get_processor_configs()
+            }
+        else:
+            processor_artifacts = {}
+
         if (
             len(processed_dataset) > 0
             and isinstance(analysis, DatasetProfilerResults)
@@ -259,6 +270,7 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         return PreviewResults(
             dataset=processed_dataset,
             analysis=analysis,
+            processor_artifacts=processor_artifacts,
             config_builder=config_builder,
         )
 

data_designer/interface/results.py

@@ -3,12 +3,15 @@
 
 from __future__ import annotations
 
+from pathlib import Path
+
 import pandas as pd
 
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.utils.visualization import WithRecordSamplerMixin
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+from data_designer.engine.dataset_builders.errors import ArtifactStorageError
 
 
 class DatasetCreationResults(WithRecordSamplerMixin):
@@ -53,3 +56,36 @@ class DatasetCreationResults(WithRecordSamplerMixin):
             A pandas DataFrame containing the full generated dataset.
         """
         return self.artifact_storage.load_dataset()
+
+    def load_processor_dataset(self, processor_name: str) -> pd.DataFrame:
+        """Load the dataset generated by a processor.
+
+        This only works for processors that write their artifacts in Parquet format.
+
+        Args:
+            processor_name: The name of the processor to load the dataset from.
+
+        Returns:
+            A pandas DataFrame containing the dataset generated by the processor.
+        """
+        try:
+            dataset = self.artifact_storage.read_parquet_files(
+                self.artifact_storage.processors_outputs_path / processor_name
+            )
+        except Exception as e:
+            raise ArtifactStorageError(f"Failed to load dataset for processor {processor_name}: {e}")
+
+        return dataset
+
+    def get_path_to_processor_artifacts(self, processor_name: str) -> Path:
+        """Get the path to the artifacts generated by a processor.
+
+        Args:
+            processor_name: The name of the processor to load the artifact from.
+
+        Returns:
+            The path to the artifacts.
+        """
+        if not self.artifact_storage.processors_outputs_path.exists():
+            raise ArtifactStorageError(f"Processor {processor_name} has no artifacts.")
+        return self.artifact_storage.processors_outputs_path / processor_name

data_designer/logging.py

@@ -6,7 +6,7 @@ import random
 import sys
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TextIO, Union
+from typing import TextIO
 
 from pythonjsonlogger import jsonlogger
 
@@ -19,7 +19,7 @@ class LoggerConfig:
 
 @dataclass
 class OutputConfig:
-    destination: Union[TextIO, Path]
+    destination: TextIO | Path
     structured: bool