refactor: update single column base class (#206)

* make properties abstract

* add private column emoji attribute

* update e2e plugin tests

* throw error if not default string

* add unit tests

* make emoji a static method

* dont need that docstring

* update unit test

Author: johnnygreco

src/data_designer/config/analysis/utils/reporting.py

@@ -16,7 +16,11 @@
 
 from data_designer.config.analysis.column_statistics import CategoricalHistogramData
 from data_designer.config.analysis.utils.errors import AnalysisReportError
-from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP, DataDesignerColumnType, get_column_display_order
+from data_designer.config.column_types import (
+    DataDesignerColumnType,
+    get_column_display_order,
+    get_column_emoji_from_type,
+)
 from data_designer.config.utils.visualization import (
     ColorPalette,
     convert_to_row_element,
@@ -101,7 +105,7 @@ def generate_analysis_report(
         displayed_column_types.add(column_type)
         column_label = column_type.replace("_", " ").title().replace("Llm", "LLM")
         table = Table(
-            title=f"{COLUMN_TYPE_EMOJI_MAP[column_type]} {column_label} Columns",
+            title=f"{get_column_emoji_from_type(column_type)} {column_label} Columns",
             **table_kws,
         )
 

src/data_designer/config/column_configs.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from abc import ABC
+from abc import ABC, abstractmethod
 from typing import Annotated, Literal
 
 from pydantic import BaseModel, Discriminator, Field, model_validator
@@ -13,7 +13,7 @@
 from data_designer.config.sampler_params import SamplerParamsT, SamplerType
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
-from data_designer.config.utils.misc import assert_valid_jinja2_template, get_prompt_template_keywords
+from data_designer.config.utils.misc import assert_valid_jinja2_template, extract_keywords_from_jinja2_template
 from data_designer.config.validator_params import ValidatorParamsT, ValidatorType
 
 
@@ -35,17 +35,22 @@ class SingleColumnConfig(ConfigBase, ABC):
     drop: bool = False
     column_type: str
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🎨"
+
     @property
+    @abstractmethod
     def required_columns(self) -> list[str]:
         """Returns a list of column names that must exist before this column can be generated.
 
         Returns:
             List of column names that this column depends on. Empty list indicates
             no dependencies. Override in subclasses to specify dependencies.
         """
-        return []
 
     @property
+    @abstractmethod
     def side_effect_columns(self) -> list[str]:
         """Returns a list of additional columns that this column will create as a side effect.
 
@@ -56,7 +61,6 @@ def side_effect_columns(self) -> list[str]:
             List of column names that this column will create as a side effect. Empty list
             indicates no side effect columns. Override in subclasses to specify side effects.
         """
-        return []
 
 
 class SamplerColumnConfig(SingleColumnConfig):
@@ -94,6 +98,18 @@ class SamplerColumnConfig(SingleColumnConfig):
     convert_to: str | None = None
     column_type: Literal["sampler"] = "sampler"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🎲"
+
+    @property
+    def required_columns(self) -> list[str]:
+        return []
+
+    @property
+    def side_effect_columns(self) -> list[str]:
+        return []
+
     @model_validator(mode="before")
     @classmethod
     def inject_sampler_type_into_params(cls, data: dict) -> dict:
@@ -150,16 +166,20 @@ class LLMTextColumnConfig(SingleColumnConfig):
     multi_modal_context: list[ImageContext] | None = None
     column_type: Literal["llm-text"] = "llm-text"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "📝"
+
     @property
     def required_columns(self) -> list[str]:
         """Get columns referenced in the prompt and system_prompt templates.
 
         Returns:
             List of unique column names referenced in Jinja2 templates.
         """
-        required_cols = list(get_prompt_template_keywords(self.prompt))
+        required_cols = list(extract_keywords_from_jinja2_template(self.prompt))
         if self.system_prompt:
-            required_cols.extend(list(get_prompt_template_keywords(self.system_prompt)))
+            required_cols.extend(list(extract_keywords_from_jinja2_template(self.system_prompt)))
         return list(set(required_cols))
 
     @property
@@ -207,6 +227,10 @@ class LLMCodeColumnConfig(LLMTextColumnConfig):
     code_lang: CodeLang
     column_type: Literal["llm-code"] = "llm-code"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "💻"
+
 
 class LLMStructuredColumnConfig(LLMTextColumnConfig):
     """Configuration for structured JSON generation columns using Large Language Models.
@@ -225,6 +249,10 @@ class LLMStructuredColumnConfig(LLMTextColumnConfig):
     output_format: dict | type[BaseModel]
     column_type: Literal["llm-structured"] = "llm-structured"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🗂️"
+
     @model_validator(mode="after")
     def validate_output_format(self) -> Self:
         """Convert Pydantic model to JSON schema if needed.
@@ -275,6 +303,10 @@ class LLMJudgeColumnConfig(LLMTextColumnConfig):
     scores: list[Score] = Field(..., min_length=1)
     column_type: Literal["llm-judge"] = "llm-judge"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "⚖️"
+
 
 class ExpressionColumnConfig(SingleColumnConfig):
     """Configuration for derived columns using Jinja2 expressions.
@@ -297,10 +329,18 @@ class ExpressionColumnConfig(SingleColumnConfig):
     dtype: Literal["int", "float", "str", "bool"] = "str"
     column_type: Literal["expression"] = "expression"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🧩"
+
     @property
     def required_columns(self) -> list[str]:
         """Returns the columns referenced in the expression template."""
-        return list(get_prompt_template_keywords(self.expr))
+        return list(extract_keywords_from_jinja2_template(self.expr))
+
+    @property
+    def side_effect_columns(self) -> list[str]:
+        return []
 
     @model_validator(mode="after")
     def assert_expression_valid_jinja(self) -> Self:
@@ -359,11 +399,19 @@ class ValidationColumnConfig(SingleColumnConfig):
     batch_size: int = Field(default=10, ge=1, description="Number of records to process in each batch")
     column_type: Literal["validation"] = "validation"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🔍"
+
     @property
     def required_columns(self) -> list[str]:
         """Returns the columns that need to be validated."""
         return self.target_columns
 
+    @property
+    def side_effect_columns(self) -> list[str]:
+        return []
+
 
 class SeedDatasetColumnConfig(SingleColumnConfig):
     """Configuration for columns sourced from seed datasets.
@@ -378,6 +426,18 @@ class SeedDatasetColumnConfig(SingleColumnConfig):
 
     column_type: Literal["seed-dataset"] = "seed-dataset"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🌱"
+
+    @property
+    def required_columns(self) -> list[str]:
+        return []
+
+    @property
+    def side_effect_columns(self) -> list[str]:
+        return []
+
 
 class EmbeddingColumnConfig(SingleColumnConfig):
     """Configuration for embedding generation columns.
@@ -395,6 +455,14 @@ class EmbeddingColumnConfig(SingleColumnConfig):
     model_alias: str
     column_type: Literal["embedding"] = "embedding"
 
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🧬"
+
     @property
     def required_columns(self) -> list[str]:
         return [self.target_column]
+
+    @property
+    def side_effect_columns(self) -> list[str]:
+        return []

src/data_designer/config/column_types.py

@@ -15,7 +15,7 @@
     SeedDatasetColumnConfig,
     ValidationColumnConfig,
 )
-from data_designer.config.errors import InvalidColumnTypeError, InvalidConfigError
+from data_designer.config.errors import InvalidConfigError
 from data_designer.config.sampler_params import SamplerType
 from data_designer.config.utils.type_helpers import (
     SAMPLER_PARAMS,
@@ -45,22 +45,6 @@
     discriminator_field_name="column_type",
 )
 
-COLUMN_TYPE_EMOJI_MAP = {
-    "general": "⚛️",  # possible analysis column type
-    DataDesignerColumnType.EXPRESSION: "🧩",
-    DataDesignerColumnType.LLM_CODE: "💻",
-    DataDesignerColumnType.LLM_JUDGE: "⚖️",
-    DataDesignerColumnType.LLM_STRUCTURED: "🗂️",
-    DataDesignerColumnType.LLM_TEXT: "📝",
-    DataDesignerColumnType.SEED_DATASET: "🌱",
-    DataDesignerColumnType.SAMPLER: "🎲",
-    DataDesignerColumnType.VALIDATION: "🔍",
-    DataDesignerColumnType.EMBEDDING: "🧬",
-}
-COLUMN_TYPE_EMOJI_MAP.update(
-    {DataDesignerColumnType(p.name): p.emoji for p in plugin_manager.get_column_generator_plugins()}
-)
-
 
 def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType, **kwargs) -> ColumnConfigT:
     """Create a Data Designer column config object from kwargs.
@@ -74,27 +58,20 @@ def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType
         Data Designer column object of the appropriate type.
     """
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    if column_type == DataDesignerColumnType.LLM_TEXT:
-        return LLMTextColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.LLM_CODE:
-        return LLMCodeColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.LLM_STRUCTURED:
-        return LLMStructuredColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.LLM_JUDGE:
-        return LLMJudgeColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.VALIDATION:
-        return ValidationColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.EXPRESSION:
-        return ExpressionColumnConfig(name=name, **kwargs)
+    config_cls = get_column_config_cls_from_type(column_type)
     if column_type == DataDesignerColumnType.SAMPLER:
-        return SamplerColumnConfig(name=name, **_resolve_sampler_kwargs(name, kwargs))
-    if column_type == DataDesignerColumnType.SEED_DATASET:
-        return SeedDatasetColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.EMBEDDING:
-        return EmbeddingColumnConfig(name=name, **kwargs)
+        kwargs = _resolve_sampler_kwargs(name, kwargs)
+    return config_cls(name=name, **kwargs)
+
+
+def get_column_config_cls_from_type(column_type: DataDesignerColumnType) -> type[ColumnConfigT]:
+    """Get the column config class for a column type."""
+    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
+    if column_type in _COLUMN_TYPE_CONFIG_CLS_MAP:
+        return _COLUMN_TYPE_CONFIG_CLS_MAP[column_type]
     if plugin := plugin_manager.get_column_generator_plugin_if_exists(column_type.value):
-        return plugin.config_cls(name=name, **kwargs)
-    raise InvalidColumnTypeError(f"🛑 {column_type} is not a valid column type.")  # pragma: no cover
+        return plugin.config_cls
+    raise InvalidConfigError(f"🛑 {column_type} is not a valid column type.")
 
 
 def get_column_display_order() -> list[DataDesignerColumnType]:
@@ -114,6 +91,12 @@ def get_column_display_order() -> list[DataDesignerColumnType]:
     return display_order
 
 
+def get_column_emoji_from_type(column_type: DataDesignerColumnType) -> str:
+    """Get the emoji for a column type."""
+    config_cls = get_column_config_cls_from_type(resolve_string_enum(column_type, DataDesignerColumnType))
+    return config_cls.get_column_emoji()
+
+
 def _resolve_sampler_kwargs(name: str, kwargs: dict) -> dict:
     if "sampler_type" not in kwargs:
         raise InvalidConfigError(f"🛑 `sampler_type` is required for sampler column '{name}'.")
@@ -142,3 +125,16 @@ def _resolve_sampler_kwargs(name: str, kwargs: dict) -> dict:
         "params": params,
         **{k: v for k, v in kwargs.items() if k not in ["sampler_type", "params"]},
     }
+
+
+_COLUMN_TYPE_CONFIG_CLS_MAP = {
+    DataDesignerColumnType.LLM_TEXT: LLMTextColumnConfig,
+    DataDesignerColumnType.LLM_CODE: LLMCodeColumnConfig,
+    DataDesignerColumnType.LLM_STRUCTURED: LLMStructuredColumnConfig,
+    DataDesignerColumnType.LLM_JUDGE: LLMJudgeColumnConfig,
+    DataDesignerColumnType.VALIDATION: ValidationColumnConfig,
+    DataDesignerColumnType.EXPRESSION: ExpressionColumnConfig,
+    DataDesignerColumnType.SAMPLER: SamplerColumnConfig,
+    DataDesignerColumnType.SEED_DATASET: SeedDatasetColumnConfig,
+    DataDesignerColumnType.EMBEDDING: EmbeddingColumnConfig,
+}

src/data_designer/config/utils/misc.py

@@ -48,8 +48,8 @@ def can_run_data_designer_locally() -> bool:
     return True
 
 
-def get_prompt_template_keywords(template: str) -> set[str]:
-    """Extract all keywords from a valid string template."""
+def extract_keywords_from_jinja2_template(template: str) -> set[str]:
+    """Extract all keywords from a valid Jinja2 template."""
     with template_error_handler():
         ast = ImmutableSandboxedEnvironment().parse(template)
         keywords = set(meta.find_undeclared_variables(ast))

src/data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -19,7 +19,7 @@
     MissingValue,
     NumericalDistribution,
 )
-from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP, DataDesignerColumnType
+from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.utils.judge_score_processing import (
     extract_judge_score_distributions,
@@ -43,8 +43,7 @@ def profile(self, column_config_with_df: ColumnConfigWithDataFrame) -> JudgeScor
         column_config, df = column_config_with_df.as_tuple()
 
         logger.info(
-            f"{COLUMN_TYPE_EMOJI_MAP[column_config.column_type]} Analyzing LLM-as-judge "
-            f"scores for column: '{column_config.name}'"
+            f"{column_config.get_column_emoji()} Analyzing LLM-as-judge scores for column: '{column_config.name}'"
         )
 
         score_summaries = {}

src/data_designer/engine/analysis/dataset_profiler.py

@@ -13,7 +13,7 @@
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import SingleColumnConfig
-from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP, ColumnConfigT
+from data_designer.config.column_types import ColumnConfigT
 from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.column_statistics import get_column_statistics_calculator
 from data_designer.engine.analysis.errors import DatasetProfilerConfigurationError
@@ -68,7 +68,7 @@ def profile_dataset(
 
         column_statistics = []
         for c in self.config.column_configs:
-            logger.info(f"  |-- {COLUMN_TYPE_EMOJI_MAP[c.column_type]} column: '{c.name}'")
+            logger.info(f"  |-- {c.get_column_emoji()} column: '{c.name}'")
             column_statistics.append(
                 get_column_statistics_calculator(c.column_type)(
                     column_config_with_df=ColumnConfigWithDataFrame(column_config=c, df=dataset)

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

@@ -95,7 +95,9 @@ def inference_parameters(self) -> BaseInferenceParams:
         return self.model_config.inference_parameters
 
     def log_pre_generation(self) -> None:
-        logger.info(f"{self.config.column_type} model configuration for generating column '{self.config.name}'")
+        logger.info(
+            f"{self.config.get_column_emoji()} {self.config.column_type} model config for column '{self.config.name}'"
+        )
         logger.info(f"  |-- model: {self.model_config.model!r}")
         logger.info(f"  |-- model alias: {self.config.model_alias!r}")
         logger.info(f"  |-- model provider: {self.get_model_provider_name(model_alias=self.config.model_alias)!r}")

src/data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -8,7 +8,7 @@
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.config.models import ModelConfig
 from data_designer.config.utils.code_lang import CodeLang
-from data_designer.config.utils.misc import get_prompt_template_keywords
+from data_designer.config.utils.misc import extract_keywords_from_jinja2_template
 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 (
@@ -56,7 +56,7 @@ def _prepare_environment(self, *, prompt_template: str | None, record: dict, pro
                 dataset_variables=list(record.keys()),
             )
         except (UserTemplateUnsupportedFiltersError, UserTemplateError) as exc:
-            template_variables = get_prompt_template_keywords(prompt_template)
+            template_variables = extract_keywords_from_jinja2_template(prompt_template)
             missing_columns = list(set(template_variables) - set(record.keys()))
 
             error_msg = (