fix: include plugin column types in display_sample_record() (#365)

* fix: include plugin column types in display_sample_record()

Replace hardcoded column type list with dynamic iteration over
get_column_display_order(), which already includes plugin-registered
types. Column types with dedicated display sections (SEED_DATASET,
IMAGE, LLM_CODE, VALIDATION, LLM_JUDGE) are excluded from the
"Generated Columns" table as before.

Also display side_effect_columns for plugin column types, matching
the existing behavior for CUSTOM columns.

Fixes #345

Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Made-with: Cursor
Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Made-with: Cursor

* fix: use is_plugin_column_type() for side-effect column display

Add is_plugin_column_type() helper to column_types.py to avoid
redundant plugin_manager.get_plugin_column_types() calls. Use it
in display_sample_record() to show side_effect_columns for plugin
column types, matching existing CUSTOM behavior.

Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Made-with: Cursor

* fix: handle string column_type in is_plugin_column_type()

Column configs store column_type as a Literal string, not a
DataDesignerColumnType enum. Accept both str and enum to avoid
AttributeError when calling .value on a plain string.

Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Made-with: Cursor

* test: improve plugin column display tests

- Move test imports to module level per project convention
- Replace plumbing-only test with one that renders a fake plugin
  column (with side-effect columns) to HTML and asserts the values
  appear in the output
- Add parametrized test for is_plugin_column_type() verifying all
  built-in types return False for both enum and string forms

Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Made-with: Cursor

* fix: preserve original display order for side-effect columns

Render primary column before side-effect columns, matching the
existing CUSTOM column behavior. Avoids introducing display
ordering discrepancies.

Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Made-with: Cursor

---------

Signed-off-by: Yev Meyer <ymeyer@nvidia.com>
Co-authored-by: Johnny Greco <jogreco@nvidia.com>

Author: 3mei

packages/data-designer-config/src/data_designer/config/column_types.py

@@ -98,6 +98,12 @@ def get_column_display_order() -> list[DataDesignerColumnType]:
     return display_order
 
 
+def is_plugin_column_type(column_type: str | DataDesignerColumnType) -> bool:
+    """Check whether a column type was registered by a plugin."""
+    type_value = column_type.value if isinstance(column_type, DataDesignerColumnType) else column_type
+    return plugin_manager.get_column_generator_plugin_if_exists(type_value) is not None
+
+
 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))

packages/data-designer-config/src/data_designer/config/utils/visualization.py

@@ -26,7 +26,7 @@
 
 import data_designer.lazy_heavy_imports as lazy
 from data_designer.config.base import ConfigBase
-from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.column_types import DataDesignerColumnType, get_column_display_order, is_plugin_column_type
 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
@@ -53,6 +53,14 @@
 console = Console()
 logger = logging.getLogger(__name__)
 
+_DEDICATED_DISPLAY_COL_TYPES = {
+    DataDesignerColumnType.SEED_DATASET,
+    DataDesignerColumnType.IMAGE,
+    DataDesignerColumnType.LLM_CODE,
+    DataDesignerColumnType.VALIDATION,
+    DataDesignerColumnType.LLM_JUDGE,
+}
+
 
 def _display_image_if_in_notebook(image_data: str, col_name: str) -> bool:
     """Display image with caption in Jupyter notebook if available.
@@ -287,14 +295,10 @@ def display_sample_record(
                 table.add_row(col_name, convert_to_row_element(record[col_name]))
         render_list.append(pad_console_element(table))
 
-    non_code_columns = (
-        config_builder.get_columns_of_type(DataDesignerColumnType.SAMPLER)
-        + 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.EMBEDDING)
-        + config_builder.get_columns_of_type(DataDesignerColumnType.CUSTOM)
-    )
+    non_code_columns = []
+    for col_type in get_column_display_order():
+        if col_type not in _DEDICATED_DISPLAY_COL_TYPES:
+            non_code_columns.extend(config_builder.get_columns_of_type(col_type))
     if len(non_code_columns) > 0:
         table = Table(title="Generated Columns", **table_kws)
         table.add_column("Name")
@@ -306,8 +310,7 @@ def display_sample_record(
                         get_truncated_list_as_string(embd) for embd in record[col.name].get("embeddings")
                     ]
                 table.add_row(col.name, convert_to_row_element(record[col.name]))
-                # Also display side_effect_columns for custom generators
-                if col.column_type == DataDesignerColumnType.CUSTOM:
+                if col.column_type == DataDesignerColumnType.CUSTOM or is_plugin_column_type(col.column_type):
                     for output_col in col.side_effect_columns:
                         if output_col in record:
                             table.add_row(output_col, convert_to_row_element(record[output_col]))

packages/data-designer-config/tests/config/test_columns.py

@@ -23,6 +23,7 @@
     DataDesignerColumnType,
     get_column_config_from_kwargs,
     get_column_display_order,
+    is_plugin_column_type,
 )
 from data_designer.config.errors import InvalidConfigError
 from data_designer.config.sampler_params import (
@@ -60,6 +61,12 @@ def test_data_designer_column_type_get_display_order():
     ]
 
 
+@pytest.mark.parametrize("col_type", list(DataDesignerColumnType))
+def test_is_plugin_column_type_false_for_builtins(col_type: DataDesignerColumnType) -> None:
+    assert is_plugin_column_type(col_type) is False
+    assert is_plugin_column_type(col_type.value) is False
+
+
 def test_sampler_column_config():
     sampler_column_config = SamplerColumnConfig(
         name="test_sampler",

packages/data-designer-config/tests/config/utils/test_visualization.py

@@ -5,10 +5,12 @@
 
 from pathlib import Path
 from typing import TYPE_CHECKING
+from unittest.mock import MagicMock, patch
 
 import pytest
 
 import data_designer.lazy_heavy_imports as lazy
+from data_designer.config.column_types import get_column_display_order
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.errors import DatasetSampleDisplayError
@@ -268,6 +270,55 @@ def test_convert_to_row_element_renders_non_scalar_types() -> None:
         table.add_row(result)
 
 
+def test_display_sample_record_includes_plugin_column_types(
+    validation_output: dict,
+    config_builder_with_validation: DataDesignerConfigBuilder,
+    tmp_path: Path,
+) -> None:
+    """Plugin columns and their side-effect columns should appear in the rendered output (fixes #345)."""
+    from types import SimpleNamespace
+
+    fake_plugin_type = "fake-plugin-type"
+    fake_col = SimpleNamespace(
+        name="plugin_output",
+        column_type=fake_plugin_type,
+        drop=False,
+        side_effect_columns=["plugin_side_effect"],
+    )
+
+    record = lazy.pd.Series(
+        {
+            "code": "print('hello world')",
+            "code_validation_result": validation_output,
+            "plugin_output": "primary plugin value",
+            "plugin_side_effect": "side effect value",
+        }
+    )
+
+    extended_order = get_column_display_order() + [fake_plugin_type]
+    original_get_columns = config_builder_with_validation.get_columns_of_type
+    mock_get_columns = MagicMock(
+        side_effect=lambda ct: [fake_col] if ct == fake_plugin_type else original_get_columns(ct)
+    )
+    config_builder_with_validation.get_columns_of_type = mock_get_columns
+
+    save_path = tmp_path / "output.html"
+    with (
+        patch("data_designer.config.utils.visualization.get_column_display_order", return_value=extended_order),
+        patch(
+            "data_designer.config.utils.visualization.is_plugin_column_type",
+            side_effect=lambda ct: ct == fake_plugin_type,
+        ),
+    ):
+        display_sample_record(record, config_builder_with_validation, save_path=save_path)
+
+    content = save_path.read_text()
+    assert "plugin_output" in content, "Plugin column name missing from rendered output"
+    assert "primary plugin value" in content, "Plugin column value missing from rendered output"
+    assert "plugin_side_effect" in content, "Plugin side-effect column name missing from rendered output"
+    assert "side effect value" in content, "Plugin side-effect column value missing from rendered output"
+
+
 def test_mixin_out_of_bounds_raises_display_error(
     validation_output: dict, config_builder_with_validation: DataDesignerConfigBuilder
 ) -> None: