data-designer 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

data_designer/engine/sampling_gen/schema_builder.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from copy import deepcopy
 
 from data_designer.config.column_configs import SamplerColumnConfig

data_designer/engine/sampling_gen/utils.py

@@ -1,9 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import numbers
+from typing import TYPE_CHECKING
+
+from data_designer.lazy_heavy_imports import np
 
-import numpy as np
+if TYPE_CHECKING:
+    import numpy as np
 
 
 def check_random_state(seed):

data_designer/engine/secret_resolver.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 import logging
 import os

data_designer/engine/validation.py

@@ -19,7 +19,7 @@ from data_designer.config.processors import ProcessorConfigT, ProcessorType
 from data_designer.config.utils.constants import RICH_CONSOLE_THEME
 from data_designer.config.utils.misc import (
     can_run_data_designer_locally,
-    get_prompt_template_keywords,
+    extract_keywords_from_jinja2_template,
 )
 from data_designer.config.validator_params import ValidatorType
 from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
@@ -302,7 +302,7 @@ def validate_schema_transform_processor(
     for processor_config in processor_configs:
         if processor_config.processor_type == ProcessorType.SCHEMA_TRANSFORM:
             for col, template in processor_config.template.items():
-                template_keywords = get_prompt_template_keywords(template)
+                template_keywords = extract_keywords_from_jinja2_template(template)
                 invalid_keywords = set(template_keywords) - all_column_names
                 if len(invalid_keywords) > 0:
                     invalid_keywords = ", ".join([f"'{k}'" for k in invalid_keywords])

data_designer/engine/validators/__init__.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.validators.base import BaseValidator, ValidationResult
 from data_designer.engine.validators.local_callable import LocalCallableValidator
 from data_designer.engine.validators.python import PythonValidator

data_designer/engine/validators/base.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import Iterator
 

data_designer/engine/validators/local_callable.py

@@ -1,14 +1,19 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import logging
+from __future__ import annotations
 
-import pandas as pd
+import logging
+from typing import TYPE_CHECKING
 
 from data_designer.config.validator_params import LocalCallableValidatorParams
 from data_designer.engine.errors import LocalCallableValidationError
 from data_designer.engine.processing.gsonschema.validators import validate
 from data_designer.engine.validators.base import BaseValidator, ValidationOutput, ValidationResult
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 

data_designer/engine/validators/python.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import ast
 import json
 import logging
@@ -8,14 +10,18 @@ import subprocess
 import tempfile
 from collections import defaultdict
 from pathlib import Path
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
-import pandas as pd
 from pydantic import BaseModel
 from ruff.__main__ import find_ruff_bin
 
 from data_designer.config.validator_params import CodeValidatorParams
 from data_designer.engine.validators.base import BaseValidator, ValidationOutput, ValidationResult
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 

data_designer/engine/validators/remote.py

@@ -1,9 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
-import httpx
 from httpx_retries import Retry, RetryTransport
 
 from data_designer.config.validator_params import RemoteValidatorParams
@@ -11,6 +13,10 @@ from data_designer.engine.errors import RemoteValidationSchemaError
 from data_designer.engine.processing.gsonschema.exceptions import JSONSchemaValidationError
 from data_designer.engine.processing.gsonschema.validators import validate
 from data_designer.engine.validators.base import BaseValidator, ValidationResult
+from data_designer.lazy_heavy_imports import httpx
+
+if TYPE_CHECKING:
+    import httpx
 
 logger = logging.getLogger(__name__)
 

data_designer/engine/validators/sql.py

@@ -1,15 +1,20 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 import re
-
-import pandas as pd
-import sqlfluff
+from typing import TYPE_CHECKING
 
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.validator_params import CodeValidatorParams
 from data_designer.engine.validators.base import BaseValidator, ValidationOutput, ValidationResult
+from data_designer.lazy_heavy_imports import pd, sqlfluff
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import sqlfluff
 
 sqlfluff_logger = logging.getLogger("sqlfluff")
 sqlfluff_logger.setLevel(logging.WARNING)

data_designer/errors.py

@@ -1,5 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 
 class DataDesignerError(Exception): ...

data_designer/essentials/__init__.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.config.default_model_settings import resolve_seed_default_model_settings
 from data_designer.config.exports import *  # noqa: F403
 from data_designer.config.run_config import RunConfig

data_designer/interface/data_designer.py

@@ -1,13 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 from pathlib import Path
-
-import pandas as pd
+from typing import TYPE_CHECKING
 
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.default_model_settings import (
     get_default_model_configs,
     get_default_model_providers_missing_api_keys,
@@ -29,14 +31,10 @@ from data_designer.config.utils.constants import (
     PREDEFINED_PROVIDERS,
 )
 from data_designer.config.utils.info import InfoType, InterfaceInfo
-from data_designer.engine.analysis.dataset_profiler import (
-    DataDesignerDatasetProfiler,
-    DatasetProfilerConfig,
-)
+from data_designer.engine.analysis.dataset_profiler import DataDesignerDatasetProfiler, DatasetProfilerConfig
 from data_designer.engine.compiler import compile_data_designer_config
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.dataset_builders.column_wise_builder import ColumnWiseDatasetBuilder
-from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.model_provider import resolve_model_provider_registry
 from data_designer.engine.resources.managed_storage import init_managed_blob_storage
 from data_designer.engine.resources.resource_provider import ResourceProvider, create_resource_provider
@@ -56,14 +54,18 @@ from data_designer.engine.secret_resolver import (
 from data_designer.interface.errors import (
     DataDesignerGenerationError,
     DataDesignerProfilingError,
-    InvalidBufferValueError,
 )
 from data_designer.interface.results import DatasetCreationResults
+from data_designer.lazy_heavy_imports import pd
 from data_designer.logging import RandomEmoji
 from data_designer.plugins.plugin import PluginType
 from data_designer.plugins.registry import PluginRegistry
 
-DEFAULT_BUFFER_SIZE = 1000
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
 
 DEFAULT_SECRET_RESOLVER = CompositeResolver([EnvironmentResolver(), PlaintextResolver()])
 
@@ -75,8 +77,6 @@ DEFAULT_SEED_READERS = [
 for plugin in PluginRegistry().get_plugins(PluginType.SEED_READER):
     DEFAULT_SEED_READERS.append(plugin.impl_cls())
 
-logger = logging.getLogger(__name__)
-
 
 class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
     """Main interface for creating datasets with Data Designer.
@@ -112,7 +112,6 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
     ):
         self._secret_resolver = secret_resolver or DEFAULT_SECRET_RESOLVER
         self._artifact_path = Path(artifact_path) if artifact_path is not None else Path.cwd() / "artifacts"
-        self._buffer_size = DEFAULT_BUFFER_SIZE
         self._run_config = RunConfig()
         self._managed_assets_path = Path(managed_assets_path or MANAGED_ASSETS_PATH)
         self._model_providers = self._resolve_model_providers(model_providers)
@@ -166,10 +165,10 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
 
         resource_provider = self._create_resource_provider(dataset_name, config_builder)
 
-        builder = self._create_dataset_builder(config_builder, resource_provider)
+        builder = self._create_dataset_builder(config_builder.build(), resource_provider)
 
         try:
-            builder.build(num_records=num_records, buffer_size=self._buffer_size)
+            builder.build(num_records=num_records)
         except Exception as e:
             raise DataDesignerGenerationError(f"🛑 Error generating dataset: {e}")
 
@@ -182,10 +181,19 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         except Exception as e:
             raise DataDesignerProfilingError(f"🛑 Error profiling dataset: {e}")
 
+        dataset_metadata = resource_provider.get_dataset_metadata()
+
+        # Update metadata with column statistics from analysis
+        if analysis:
+            builder.artifact_storage.update_metadata(
+                {"column_statistics": [stat.model_dump(mode="json") for stat in analysis.column_statistics]}
+            )
+
         return DatasetCreationResults(
             artifact_storage=builder.artifact_storage,
             analysis=analysis,
             config_builder=config_builder,
+            dataset_metadata=dataset_metadata,
         )
 
     def preview(
@@ -211,7 +219,7 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         logger.info(f"{RandomEmoji.previewing()} Preview generation in progress")
 
         resource_provider = self._create_resource_provider("preview-dataset", config_builder)
-        builder = self._create_dataset_builder(config_builder, resource_provider)
+        builder = self._create_dataset_builder(config_builder.build(), resource_provider)
 
         try:
             raw_dataset = builder.build_preview(num_records=num_records)
@@ -249,11 +257,15 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         ):
             logger.info(f"{RandomEmoji.success()} Preview complete!")
 
+        # Create dataset metadata from the resource provider
+        dataset_metadata = resource_provider.get_dataset_metadata()
+
         return PreviewResults(
             dataset=processed_dataset,
             analysis=analysis,
             processor_artifacts=processor_artifacts,
             config_builder=config_builder,
+            dataset_metadata=dataset_metadata,
         )
 
     def validate(self, config_builder: DataDesignerConfigBuilder) -> None:
@@ -271,7 +283,7 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
             InvalidConfigError: If the configuration is invalid.
         """
         resource_provider = self._create_resource_provider("validate-configuration", config_builder)
-        compile_data_designer_config(config_builder, resource_provider)
+        compile_data_designer_config(config_builder.build(), resource_provider)
 
     def get_default_model_configs(self) -> list[ModelConfig]:
         """Get the default model configurations.
@@ -300,34 +312,22 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         """
         return self._secret_resolver
 
-    def set_buffer_size(self, buffer_size: int) -> None:
-        """Set the buffer size for dataset generation.
-
-        The buffer size controls how many records are processed in memory at once
-        during dataset generation using the `create` method. The default value is
-        set to the constant `DEFAULT_BUFFER_SIZE` defined in the data_designer module.
-
-        Args:
-            buffer_size: Number of records to process in each buffer.
-
-        Raises:
-            InvalidBufferValueError: If buffer size is less than or equal to 0.
-        """
-        if buffer_size <= 0:
-            raise InvalidBufferValueError("Buffer size must be greater than 0.")
-        self._buffer_size = buffer_size
-
     def set_run_config(self, run_config: RunConfig) -> None:
         """Set the runtime configuration for dataset generation.
 
         Args:
             run_config: A RunConfig instance containing runtime settings such as
-                early shutdown behavior. Import RunConfig from data_designer.essentials.
+                early shutdown behavior and batch sizing via `buffer_size`. Import RunConfig from
+                data_designer.essentials.
 
         Example:
             >>> from data_designer.essentials import DataDesigner, RunConfig
             >>> dd = DataDesigner()
             >>> dd.set_run_config(RunConfig(disable_early_shutdown=True))
+
+        Notes:
+            When `disable_early_shutdown=True`, DataDesigner will never terminate generation early
+            due to error-rate thresholds. Errors are still tracked for reporting.
         """
         self._run_config = run_config
 
@@ -348,14 +348,11 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
 
     def _create_dataset_builder(
         self,
-        config_builder: DataDesignerConfigBuilder,
+        data_designer_config: DataDesignerConfig,
         resource_provider: ResourceProvider,
     ) -> ColumnWiseDatasetBuilder:
-        config = compile_data_designer_config(config_builder, resource_provider)
-
         return ColumnWiseDatasetBuilder(
-            column_configs=compile_dataset_builder_column_configs(config),
-            processor_configs=config.processors or [],
+            data_designer_config=data_designer_config,
             resource_provider=resource_provider,
         )
 

data_designer/interface/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.errors import DataDesignerError
 
 

data_designer/interface/results.py

@@ -4,14 +4,18 @@
 from __future__ import annotations
 
 from pathlib import Path
-
-import pandas as pd
+from typing import TYPE_CHECKING
 
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.dataset_metadata import DatasetMetadata
 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
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class DatasetCreationResults(WithRecordSamplerMixin):
@@ -28,6 +32,7 @@ class DatasetCreationResults(WithRecordSamplerMixin):
         artifact_storage: ArtifactStorage,
         analysis: DatasetProfilerResults,
         config_builder: DataDesignerConfigBuilder,
+        dataset_metadata: DatasetMetadata,
     ):
         """Creates a new instance with results based on a dataset creation run.
 
@@ -35,10 +40,12 @@ class DatasetCreationResults(WithRecordSamplerMixin):
             artifact_storage: Storage manager for accessing generated artifacts.
             analysis: Profiling results for the generated dataset.
             config_builder: Configuration builder used to create the dataset.
+            dataset_metadata: Metadata about the generated dataset (e.g., seed column names).
         """
         self.artifact_storage = artifact_storage
         self._analysis = analysis
         self._config_builder = config_builder
+        self.dataset_metadata = dataset_metadata
 
     def load_analysis(self) -> DatasetProfilerResults:
         """Load the profiling analysis results for the generated dataset.

data_designer/lazy_heavy_imports.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Lazy imports facade for heavy third-party dependencies.
+
+This module provides a centralized facade that lazily imports heavy dependencies
+only when accessed, significantly improving import performance.
+
+Usage:
+    from data_designer.lazy_heavy_imports import pd, np, faker, litellm
+
+    df = pd.DataFrame(...)
+    arr = np.array([1, 2, 3])
+    fake = faker.Faker()
+"""
+
+from __future__ import annotations
+
+import importlib
+
+# Mapping of lazy import names to their actual module paths
+_LAZY_IMPORTS = {
+    "pd": "pandas",
+    "np": "numpy",
+    "pq": "pyarrow.parquet",
+    "pa": "pyarrow",
+    "faker": "faker",
+    "litellm": "litellm",
+    "sqlfluff": "sqlfluff",
+    "httpx": "httpx",
+    "duckdb": "duckdb",
+    "nx": "networkx",
+    "scipy": "scipy",
+    "jsonschema": "jsonschema",
+}
+
+
+def __getattr__(name: str) -> object:
+    """Lazily import heavy third-party dependencies when accessed.
+
+    This allows fast imports of data_designer while deferring loading of heavy
+    libraries until they're actually needed.
+    """
+    if name in _LAZY_IMPORTS:
+        module_name = _LAZY_IMPORTS[name]
+        return importlib.import_module(module_name)
+
+    raise AttributeError(f"module 'data_designer.lazy_heavy_imports' has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    """Return list of available lazy imports."""
+    return list(_LAZY_IMPORTS.keys())

data_designer/logging.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 import random
 import sys

data_designer/plugins/__init__.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.plugins.plugin import Plugin, PluginType
 
 __all__ = ["Plugin", "PluginType"]

data_designer/plugins/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.errors import DataDesignerError
 
 

data_designer/plugins/plugin.py

@@ -70,7 +70,6 @@ class Plugin(BaseModel):
         ..., description="The fully-qualified name o the config class object, e.g. 'my_plugin.config.MyConfig'"
     )
     plugin_type: PluginType = Field(..., description="The type of plugin")
-    emoji: str = Field(default="🔌", description="The emoji to use in logs related to the plugin")
 
     @property
     def config_type_as_class_name(self) -> str:

data_designer/plugins/registry.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 import os
 import threading

data_designer/plugins/testing/__init__.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.plugins.testing.utils import assert_valid_plugin
 
 __all__ = [

data_designer/plugins/testing/stubs.py

@@ -1,11 +1,13 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from typing import Literal
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import SingleColumnConfig
-from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorCellByCell
 from data_designer.plugins.plugin import Plugin, PluginType
 
 MODULE_NAME = __name__
@@ -18,15 +20,11 @@ class ValidTestConfig(SingleColumnConfig):
     name: str
 
 
-class ValidTestTask(ConfigurableTask[ValidTestConfig]):
+class ValidTestTask(ColumnGeneratorCellByCell[ValidTestConfig]):
     """Valid task for testing plugin creation."""
 
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="test_generator",
-            description="Test generator",
-        )
+    def generate(self, data: dict) -> dict:
+        return data
 
 
 class ConfigWithoutDiscriminator(ConfigBase):
@@ -53,22 +51,14 @@ class StubPluginConfigB(SingleColumnConfig):
     column_type: Literal["test-plugin-b"] = "test-plugin-b"
 
 
-class StubPluginTaskA(ConfigurableTask[StubPluginConfigA]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="test_plugin_a",
-            description="Test plugin A",
-        )
+class StubPluginTaskA(ColumnGeneratorCellByCell[StubPluginConfigA]):
+    def generate(self, data: dict) -> dict:
+        return data
 
 
-class StubPluginTaskB(ConfigurableTask[StubPluginConfigB]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="test_plugin_b",
-            description="Test plugin B",
-        )
+class StubPluginTaskB(ColumnGeneratorCellByCell[StubPluginConfigB]):
+    def generate(self, data: dict) -> dict:
+        return data
 
 
 # Stub plugins requiring different combinations of resources
@@ -86,31 +76,19 @@ class StubPluginConfigBlobsAndSeeds(SingleColumnConfig):
     column_type: Literal["test-plugin-blobs-and-seeds"] = "test-plugin-blobs-and-seeds"
 
 
-class StubPluginTaskModels(ConfigurableTask[StubPluginConfigModels]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="test_plugin_models",
-            description="Test plugin requiring models",
-        )
+class StubPluginTaskModels(ColumnGeneratorCellByCell[StubPluginConfigModels]):
+    def generate(self, data: dict) -> dict:
+        return data
 
 
-class StubPluginTaskModelsAndBlobs(ConfigurableTask[StubPluginConfigModelsAndBlobs]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="test_plugin_models_and_blobs",
-            description="Test plugin requiring models and blobs",
-        )
+class StubPluginTaskModelsAndBlobs(ColumnGeneratorCellByCell[StubPluginConfigModelsAndBlobs]):
+    def generate(self, data: dict) -> dict:
+        return data
 
 
-class StubPluginTaskBlobsAndSeeds(ConfigurableTask[StubPluginConfigBlobsAndSeeds]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="test_plugin_blobs_and_seeds",
-            description="Test plugin requiring blobs and seeds",
-        )
+class StubPluginTaskBlobsAndSeeds(ColumnGeneratorCellByCell[StubPluginConfigBlobsAndSeeds]):
+    def generate(self, data: dict) -> dict:
+        return data
 
 
 plugin_none = Plugin(

data_designer/plugins/testing/utils.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.config.base import ConfigBase
 from data_designer.engine.configurable_task import ConfigurableTask
 from data_designer.engine.resources.seed_reader import SeedReader