scaffold

Author: andreatgretel

examples/.gitignore

@@ -0,0 +1 @@
+**/artifacts

examples/custom_column.py

@@ -0,0 +1,72 @@
+import pandas as pd
+
+from data_designer.essentials import (
+    CustomColumnConfig,
+    DataDesigner,
+    DataDesignerConfigBuilder,
+    InferenceParameters,
+    LoggingConfig,
+    ModelConfig,
+    configure_logging,
+)
+
+configure_logging(LoggingConfig.debug())
+
+# Initialize NDD and add columns
+MODEL_ALIAS = "nano"
+SYSTEM_PROMPT = "/no_think"
+
+model_configs = [
+    ModelConfig(
+        alias="nano",
+        model="nvidia/nvidia-nemotron-nano-9b-v2",
+        inference_parameters=InferenceParameters(
+            temperature=0.5,
+            top_p=1.0,
+            max_tokens=1024,
+            max_parallel_requests=4,
+        ),
+        provider="nvidia",
+    )
+]
+
+builder = DataDesignerConfigBuilder(model_configs=model_configs)
+
+builder.add_column(
+    name="topic",
+    column_type="sampler",
+    sampler_type="category",
+    params={
+        "values": [
+            "healthcare",
+            "finance",
+            "technology",
+        ]
+    }
+)
+
+builder.add_column(
+    name="text",
+    column_type="llm-text",
+    model_alias=MODEL_ALIAS,
+    prompt="Write me a paragraph about {{ topic }}.",
+    system_prompt=SYSTEM_PROMPT,
+)
+
+def generator_function(df: pd.DataFrame) -> pd.DataFrame:
+    df["length_frac"] = df["text"].apply(lambda x: len(x) / 1000)
+    return df
+
+builder.add_column(
+    CustomColumnConfig(
+        name="length_frac",
+        generator_function=generator_function,
+    )
+)
+
+# Generate dataset
+dd = DataDesigner(artifact_path="./artifacts")
+dd_preview = dd.preview(builder, num_records=10)
+dd_preview.display_sample_record()
+
+dd.create(builder, num_records=20)

src/data_designer/config/columns.py

@@ -3,9 +3,10 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Literal, Optional, Type, Union
+from typing import Callable, Literal, Optional, Type, Union
 
-from pydantic import BaseModel, Field, model_validator
+import pandas as pd
+from pydantic import BaseModel, Field, field_serializer, model_validator
 from typing_extensions import Self, TypeAlias
 
 from .base import ConfigBase
@@ -28,6 +29,7 @@ class DataDesignerColumnType(str, Enum):
     EXPRESSION = "expression"
     VALIDATION = "validation"
     SEED_DATASET = "seed-dataset"
+    CUSTOM = "custom"
 
     @staticmethod
     def get_display_order() -> list[Self]:
@@ -40,6 +42,7 @@ def get_display_order() -> list[Self]:
             DataDesignerColumnType.LLM_JUDGE,
             DataDesignerColumnType.VALIDATION,
             DataDesignerColumnType.EXPRESSION,
+            DataDesignerColumnType.CUSTOM,
         ]
 
     @property
@@ -55,6 +58,7 @@ def is_dag_column_type(self) -> bool:
             self.LLM_STRUCTURED,
             self.LLM_TEXT,
             self.VALIDATION,
+            self.CUSTOM,
         ]
 
 
@@ -195,6 +199,18 @@ def column_type(self) -> DataDesignerColumnType:
         return DataDesignerColumnType.SEED_DATASET
 
 
+class CustomColumnConfig(SingleColumnConfig):
+    generator_function: Callable[[pd.DataFrame], pd.DataFrame]
+
+    @property
+    def column_type(self) -> DataDesignerColumnType:
+        return DataDesignerColumnType.CUSTOM
+
+    @field_serializer("generator_function")
+    def serialize_generator_function(self, v: Callable[[pd.DataFrame], pd.DataFrame]) -> str:
+        return v.__name__
+
+
 COLUMN_TYPE_EMOJI_MAP = {
     "general": "⚛️",  # possible analysis column type
     DataDesignerColumnType.EXPRESSION: "🧩",
@@ -205,6 +221,7 @@ def column_type(self) -> DataDesignerColumnType:
     DataDesignerColumnType.SEED_DATASET: "🌱",
     DataDesignerColumnType.SAMPLER: "🎲",
     DataDesignerColumnType.VALIDATION: "🔍",
+    DataDesignerColumnType.CUSTOM: "🛠️",
 }
 
 
@@ -217,6 +234,7 @@ def column_type(self) -> DataDesignerColumnType:
     SamplerColumnConfig,
     SeedDatasetColumnConfig,
     ValidationColumnConfig,
+    CustomColumnConfig,
 ]
 
 
@@ -248,6 +266,8 @@ def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType
         return SamplerColumnConfig(name=name, **_resolve_sampler_kwargs(name, kwargs))
     elif column_type == DataDesignerColumnType.SEED_DATASET:
         return SeedDatasetColumnConfig(name=name, **kwargs)
+    elif column_type == DataDesignerColumnType.CUSTOM:
+        return CustomColumnConfig(name=name, **kwargs)
     raise InvalidColumnTypeError(f"🛑 {column_type} is not a valid column type.")  # pragma: no cover
 
 

src/data_designer/config/utils/visualization.py

@@ -160,6 +160,7 @@ def display_sample_record(
         + config_builder.get_columns_of_type(DataDesignerColumnType.EXPRESSION)
         + config_builder.get_columns_of_type(DataDesignerColumnType.LLM_TEXT)
         + config_builder.get_columns_of_type(DataDesignerColumnType.LLM_STRUCTURED)
+        + config_builder.get_columns_of_type(DataDesignerColumnType.CUSTOM)
     )
     if len(non_code_columns) > 0:
         table = Table(title="Generated Columns", **table_kws)

src/data_designer/engine/column_generators/generators/custom.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import pandas as pd
+
+from data_designer.config.columns import CustomColumnConfig
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+)
+from data_designer.engine.errors import DataDesignerRuntimeError
+
+
+logger = logging.getLogger(__name__)
+
+
+class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="custom",
+            description="Generate a custom column.",
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            required_resources=None,
+        )
+
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        logger.info(f"🛠️ Generating custom column {self.config.name!r} with {len(data)} records")
+        logger.info(f"  |-- generator function: {self.config.generator_function.__name__}")
+
+        try:
+            result = self.config.generator_function(data)
+        except Exception as e:
+            raise DataDesignerRuntimeError(f"Error generating custom column {self.config.name!r}: {e}")
+
+        return result

src/data_designer/engine/column_generators/registry.py

@@ -3,6 +3,7 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.columns import (
+    CustomColumnConfig,
     DataDesignerColumnType,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
@@ -12,6 +13,7 @@
     ValidationColumnConfig,
 )
 from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.custom import CustomColumnGenerator
 from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
 from data_designer.engine.column_generators.generators.llm_generators import (
     LLMCodeCellGenerator,
@@ -39,6 +41,7 @@ def create_default_column_generator_registry() -> ColumnGeneratorRegistry:
     registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig, False)
     registry.register(DataDesignerColumnType.EXPRESSION, ExpressionColumnGenerator, ExpressionColumnConfig, False)
     registry.register(DataDesignerColumnType.SAMPLER, SamplerColumnGenerator, SamplerMultiColumnConfig, False)
+    registry.register(DataDesignerColumnType.CUSTOM, CustomColumnGenerator, CustomColumnConfig, False)
     registry.register(
         DataDesignerColumnType.SEED_DATASET,
         SeedDatasetColumnGenerator,

src/data_designer/essentials/__init__.py

@@ -3,6 +3,7 @@
 
 from ..config.analysis.column_profilers import JudgeScoreProfilerConfig
 from ..config.columns import (
+    CustomColumnConfig,
     DataDesignerColumnType,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,