data-designer 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

data_designer/config/utils/visualization.py

@@ -3,11 +3,11 @@
 
 from __future__ import annotations
 
+import json
+import os
 from collections import OrderedDict
 from enum import Enum
 from functools import cached_property
-import json
-import os
 from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
@@ -21,16 +21,16 @@ from rich.syntax import Syntax
 from rich.table import Table
 from rich.text import Text
 
-from ..base import ConfigBase
-from ..column_types import DataDesignerColumnType
-from ..models import ModelConfig, ModelProvider
-from ..sampler_params import SamplerType
-from .code_lang import code_lang_to_syntax_lexer
-from .constants import NVIDIA_API_KEY_ENV_VAR_NAME, OPENAI_API_KEY_ENV_VAR_NAME
-from .errors import DatasetSampleDisplayError
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.models import ModelConfig, ModelProvider
+from data_designer.config.sampler_params import SamplerType
+from data_designer.config.utils.code_lang import code_lang_to_syntax_lexer
+from data_designer.config.utils.constants import NVIDIA_API_KEY_ENV_VAR_NAME, OPENAI_API_KEY_ENV_VAR_NAME
+from data_designer.config.utils.errors import DatasetSampleDisplayError
 
 if TYPE_CHECKING:
-    from ..config_builder import DataDesignerConfigBuilder
+    from data_designer.config.config_builder import DataDesignerConfigBuilder
 
 
 console = Console()
@@ -72,6 +72,9 @@ class WithRecordSamplerMixin:
         else:
             raise DatasetSampleDisplayError("No valid dataset found in results object.")
 
+    def _has_processor_artifacts(self) -> bool:
+        return hasattr(self, "processor_artifacts") and self.processor_artifacts is not None
+
     def display_sample_record(
         self,
         index: Optional[int] = None,
@@ -79,6 +82,7 @@ class WithRecordSamplerMixin:
         hide_seed_columns: bool = False,
         syntax_highlighting_theme: str = "dracula",
         background_color: Optional[str] = None,
+        processors_to_display: Optional[list[str]] = None,
     ) -> None:
         """Display a sample record from the Data Designer dataset preview.
 
@@ -90,6 +94,7 @@ class WithRecordSamplerMixin:
                 documentation from `rich` for information about available themes.
             background_color: Background color to use for the record. See the `Syntax`
                 documentation from `rich` for information about available background colors.
+            processors_to_display: List of processors to display the artifacts for. If None, all processors will be displayed.
         """
         i = index or self._display_cycle_index
 
@@ -99,8 +104,25 @@ class WithRecordSamplerMixin:
         except IndexError:
             raise DatasetSampleDisplayError(f"Index {i} is out of bounds for dataset of length {num_records}.")
 
+        processor_data_to_display = None
+        if self._has_processor_artifacts() and len(self.processor_artifacts) > 0:
+            if processors_to_display is None:
+                processors_to_display = list(self.processor_artifacts.keys())
+
+            if len(processors_to_display) > 0:
+                processor_data_to_display = {}
+                for processor in processors_to_display:
+                    if (
+                        isinstance(self.processor_artifacts[processor], list)
+                        and len(self.processor_artifacts[processor]) == num_records
+                    ):
+                        processor_data_to_display[processor] = self.processor_artifacts[processor][i]
+                    else:
+                        processor_data_to_display[processor] = self.processor_artifacts[processor]
+
         display_sample_record(
             record=record,
+            processor_data_to_display=processor_data_to_display,
             config_builder=self._config_builder,
             background_color=background_color,
             syntax_highlighting_theme=syntax_highlighting_theme,
@@ -134,6 +156,7 @@ def create_rich_histogram_table(
 def display_sample_record(
     record: Union[dict, pd.Series, pd.DataFrame],
     config_builder: DataDesignerConfigBuilder,
+    processor_data_to_display: Optional[dict[str, Union[list[str], str]]] = None,
     background_color: Optional[str] = None,
     syntax_highlighting_theme: str = "dracula",
     record_index: Optional[int] = None,
@@ -230,6 +253,15 @@ def display_sample_record(
             table.add_row(*row)
             render_list.append(pad_console_element(table, (1, 0, 1, 0)))
 
+    if processor_data_to_display and len(processor_data_to_display) > 0:
+        for processor_name, processor_data in processor_data_to_display.items():
+            table = Table(title=f"Processor Outputs: {processor_name}", **table_kws)
+            table.add_column("Name")
+            table.add_column("Value")
+            for col, value in processor_data.items():
+                table.add_row(col, convert_to_row_element(value))
+        render_list.append(pad_console_element(table, (1, 0, 1, 0)))
+
     if record_index is not None:
         index_label = Text(f"[index: {record_index}]", justify="center")
         render_list.append(index_label)

data_designer/config/validator_params.py

@@ -7,8 +7,8 @@ from typing import Any, Optional, Union
 from pydantic import Field, field_serializer, model_validator
 from typing_extensions import Self, TypeAlias
 
-from .base import ConfigBase
-from .utils.code_lang import SQL_DIALECTS, CodeLang
+from data_designer.config.base import ConfigBase
+from data_designer.config.utils.code_lang import SQL_DIALECTS, CodeLang
 
 SUPPORTED_CODE_LANGUAGES = {CodeLang.PYTHON, *SQL_DIALECTS}
 

data_designer/engine/analysis/column_profilers/base.py

@@ -3,8 +3,8 @@
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
 import logging
+from abc import ABC, abstractmethod
 
 import pandas as pd
 import pyarrow as pa

data_designer/engine/analysis/dataset_profiler.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 from collections.abc import Sequence
 from functools import cached_property
-import logging
 
 import pandas as pd
 from pydantic import Field, field_validator

data_designer/engine/analysis/utils/judge_score_processing.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections import defaultdict
 import logging
+from collections import defaultdict
 from typing import Any, Optional, Union
 
 import pandas as pd

data_designer/engine/column_generators/generators/samplers.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from functools import partial
 import logging
 import random
+from functools import partial
 from typing import Callable
 
 import pandas as pd

data_designer/engine/dataset_builders/artifact_storage.py

@@ -1,12 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from datetime import datetime
-from functools import cached_property
 import json
 import logging
-from pathlib import Path
 import shutil
+from datetime import datetime
+from functools import cached_property
+from pathlib import Path
 from typing import Union
 
 import pandas as pd
@@ -25,6 +25,7 @@ class BatchStage(StrEnum):
     PARTIAL_RESULT = "partial_results_path"
     FINAL_RESULT = "final_dataset_path"
     DROPPED_COLUMNS = "dropped_columns_dataset_path"
+    PROCESSORS_OUTPUTS = "processors_outputs_path"
 
 
 class ArtifactStorage(BaseModel):
@@ -33,6 +34,7 @@ class ArtifactStorage(BaseModel):
     final_dataset_folder_name: str = "parquet-files"
     partial_results_folder_name: str = "tmp-partial-parquet-files"
     dropped_columns_folder_name: str = "dropped-columns-parquet-files"
+    processors_outputs_folder_name: str = "processors-files"
 
     @property
     def artifact_path_exists(self) -> bool:
@@ -70,6 +72,10 @@ class ArtifactStorage(BaseModel):
     def partial_results_path(self) -> Path:
         return self.base_dataset_path / self.partial_results_folder_name
 
+    @property
+    def processors_outputs_path(self) -> Path:
+        return self.base_dataset_path / self.processors_outputs_folder_name
+
     @field_validator("artifact_path")
     def validate_artifact_path(cls, v: Union[Path, str]) -> Path:
         v = Path(v)
@@ -84,6 +90,7 @@ class ArtifactStorage(BaseModel):
             self.final_dataset_folder_name,
             self.partial_results_folder_name,
             self.dropped_columns_folder_name,
+            self.processors_outputs_folder_name,
         ]
 
         for name in folder_names:
@@ -169,9 +176,10 @@ class ArtifactStorage(BaseModel):
         batch_number: int,
         dataframe: pd.DataFrame,
         batch_stage: BatchStage,
+        subfolder: str | None = None,
     ) -> Path:
         file_path = self.create_batch_file_path(batch_number, batch_stage=batch_stage)
-        self.write_parquet_file(file_path.name, dataframe, batch_stage)
+        self.write_parquet_file(file_path.name, dataframe, batch_stage, subfolder=subfolder)
         return file_path
 
     def write_parquet_file(
@@ -179,9 +187,11 @@ class ArtifactStorage(BaseModel):
         parquet_file_name: str,
         dataframe: pd.DataFrame,
         batch_stage: BatchStage,
+        subfolder: str | None = None,
     ) -> Path:
-        self.mkdir_if_needed(self._get_stage_path(batch_stage))
-        file_path = self._get_stage_path(batch_stage) / parquet_file_name
+        subfolder = subfolder or ""
+        self.mkdir_if_needed(self._get_stage_path(batch_stage) / subfolder)
+        file_path = self._get_stage_path(batch_stage) / subfolder / parquet_file_name
         dataframe.to_parquet(file_path, index=False)
         return file_path
 

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -4,8 +4,8 @@
 import functools
 import json
 import logging
-from pathlib import Path
 import time
+from pathlib import Path
 from typing import Callable
 
 import pandas as pd
@@ -171,6 +171,8 @@ class ColumnWiseDatasetBuilder:
         max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
         if isinstance(generator, WithLLMGeneration):
             max_workers = generator.inference_parameters.max_parallel_requests
+        elif hasattr(generator.config, "max_parallel_requests"):
+            max_workers = generator.config.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)
 
     def _run_full_column_generator(self, generator: ColumnGenerator) -> None:
@@ -244,6 +246,7 @@ class ColumnWiseDatasetBuilder:
             processors[BuildStage.POST_BATCH].append(  # as post-batch by default
                 DropColumnsProcessor(
                     config=DropColumnsProcessorConfig(
+                        name="default_drop_columns_processor",
                         column_names=columns_to_drop,
                         build_stage=BuildStage.POST_BATCH,
                     ),

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -3,10 +3,10 @@
 
 from __future__ import annotations
 
-from concurrent.futures import Future, ThreadPoolExecutor
 import contextvars
 import json
 import logging
+from concurrent.futures import Future, ThreadPoolExecutor
 from threading import Lock, Semaphore
 from typing import Any, Optional, Protocol
 

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from pathlib import Path
 import shutil
+from pathlib import Path
 from typing import Callable, Container, Iterator
 
 import pandas as pd

data_designer/engine/errors.py

@@ -3,7 +3,7 @@
 
 from pydantic import BaseModel, Field
 
-from ..errors import DataDesignerError
+from data_designer.errors import DataDesignerError
 
 
 class DataDesignerRuntimeError(DataDesignerError): ...

data_designer/engine/models/errors.py

@@ -3,9 +3,9 @@
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Callable
 from functools import wraps
-import logging
 from typing import Any
 
 from litellm.exceptions import (

data_designer/engine/models/facade.py

@@ -3,9 +3,9 @@
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Callable
 from copy import deepcopy
-import logging
 from typing import Any
 
 from litellm.types.router import DeploymentTypedDict, LiteLLM_Params

data_designer/engine/models/parsers/parser.py

@@ -4,12 +4,12 @@
 from functools import reduce
 from typing import Optional
 
+import marko
 from lxml import etree
 from lxml.etree import _Element
-import marko
 
-from data_designer.engine.models.parsers.postprocessors import merge_text_blocks
 import data_designer.engine.models.parsers.tag_parsers as tp
+from data_designer.engine.models.parsers.postprocessors import merge_text_blocks
 from data_designer.engine.models.parsers.types import (
     LLMStructuredResponse,
     PostProcessor,

data_designer/engine/models/recipes/response_recipes.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections.abc import Callable
 import json
+from collections.abc import Callable
 
 from pydantic import BaseModel
 

data_designer/engine/processing/ginja/environment.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import re
 from collections.abc import Callable
 from functools import partial, wraps
-import re
 from typing import Any
 
 from jinja2 import meta

data_designer/engine/processing/gsonschema/validators.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from copy import deepcopy
 import logging
+from copy import deepcopy
 from typing import Any, overload
 
 from jsonschema import Draft202012Validator, ValidationError, validators

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/resources/managed_dataset_repository.py

@@ -1,14 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from functools import cached_property
 import logging
-from pathlib import Path
 import tempfile
 import threading
 import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from functools import cached_property
+from pathlib import Path
 from typing import Any
 
 import duckdb

data_designer/engine/resources/managed_storage.py

@@ -1,10 +1,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
 from contextlib import contextmanager
-import logging
 from pathlib import Path
 from typing import IO
 

data_designer/engine/sampling_gen/constraints.py

@@ -5,8 +5,8 @@ from abc import ABC, abstractmethod
 from typing import Type
 
 import numpy as np
-from numpy.typing import NDArray
 import pandas as pd
+from numpy.typing import NDArray
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.sampler_constraints import (

data_designer/engine/sampling_gen/data_sources/base.py

@@ -5,8 +5,8 @@ from abc import ABC, abstractmethod
 from typing import Any, Generic, Optional, Type, TypeVar, Union
 
 import numpy as np
-from numpy.typing import NDArray
 import pandas as pd
+from numpy.typing import NDArray
 from scipy import stats
 
 from data_designer.config.sampler_params import SamplerParamsT

data_designer/engine/sampling_gen/entities/email_address_utils.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from datetime import date
 import random
 import re
+from datetime import date
 
 import anyascii
 

data_designer/engine/sampling_gen/entities/national_id_utils.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from datetime import date
 import random
+from datetime import date
 
 SSN_RANDOMIZATION_DATE = date(2011, 6, 25)
 

data_designer/engine/sampling_gen/entities/person.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from datetime import date, timedelta
 import random
+from datetime import date, timedelta
 from typing import Any, Literal, TypeAlias
 
 from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS

data_designer/engine/sampling_gen/entities/phone_number.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from pathlib import Path
 import random
+from pathlib import Path
 from typing import Optional
 
 import pandas as pd

data_designer/engine/sampling_gen/people_gen.py

@@ -3,15 +3,15 @@
 
 from __future__ import annotations
 
+import random
+import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from copy import deepcopy
-import random
 from typing import TYPE_CHECKING, Any, Union
-import uuid
 
-from faker import Faker
 import pandas as pd
+from faker import Faker
 
 from data_designer.config.utils.constants import AVAILABLE_LOCALES, DEFAULT_AGE_RANGE
 from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator

data_designer/engine/secret_resolver.py

@@ -1,10 +1,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections.abc import Sequence
 import json
 import logging
 import os
+from collections.abc import Sequence
 from pathlib import Path
 from typing import Protocol
 

data_designer/engine/validators/python.py

@@ -2,12 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import ast
-from collections import defaultdict
 import logging
-from pathlib import Path
 import re
 import subprocess
 import tempfile
+from collections import defaultdict
+from pathlib import Path
 from uuid import uuid4
 
 import pandas as pd