Refactor tests to support group-atoms and improve converter handling

- Updated `test_runs_api.py` to utilize group-atoms for input and output columns in model sessions, ensuring converters are correctly applied and persisted.
- Added new tests for predicting with session converters and validating the behavior of group atoms in input/output columns.
- Enhanced `test_session_preprocessing_job.py` to verify that preprocessing correctly saves separate x and y files, and resolves group atoms per fold in cross-validation.
- Introduced tests for handling output columns as group atoms, ensuring proper error handling when such configurations are attempted.
- Improved `test_dashai_dataset.py` to ensure that column selection correctly prunes saved type metadata, preventing stale schema information from persisting after dataset modifications.

Author: Creylay

DashAI/back/api/api_v1/endpoints/explainers.py

@@ -1023,10 +1023,22 @@ async def validate_dataset(
                 detail="Internal database error",
             ) from e
 
+    from DashAI.back.job.base_job import JobError
+    from DashAI.back.job.session_preprocessing_job import (
+        get_real_input_output_columns,
+    )
+
     validation_response = {}
-    input_columns = model_session.input_columns
-    output_columns = model_session.output_columns
-    required_columns = input_columns + output_columns
+    try:
+        real_input_columns, real_output_columns = get_real_input_output_columns(
+            model_session
+        )
+    except JobError as e:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        ) from e
+    required_columns = real_input_columns + real_output_columns
 
     instances_columns = list(instances.features)
 
@@ -1122,14 +1134,29 @@ async def valid_datasets(
                 detail="Internal database error",
             ) from e
 
+    from DashAI.back.job.base_job import JobError
+    from DashAI.back.job.session_preprocessing_job import (
+        get_real_input_output_columns,
+    )
+
+    try:
+        real_input_columns, real_output_columns = get_real_input_output_columns(
+            model_session
+        )
+    except JobError as e:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        ) from e
+
     # Fallback only: a converter that adds/renames input columns
     # (BagOfWords' `bow_<word>`, PCA's `pca0`/`pca1`) means
     # `model_session.input_columns` only exist in the session's
     # preprocessed data, never in a raw dataset — every candidate used to
     # fail this check. Replaced below with the training dataset's own raw
     # schema whenever it's readable; kept here only for the rare case that
     # read fails, so this doesn't silently accept every dataset.
-    required_columns = model_session.input_columns + model_session.output_columns
+    required_columns = real_input_columns + real_output_columns
 
     training_types = {}
     training_spec = {}
@@ -1144,8 +1171,8 @@ async def valid_datasets(
 
     if training_spec:
         required_columns = [
-            col for col in training_spec if col not in model_session.output_columns
-        ] + model_session.output_columns
+            col for col in training_spec if col not in real_output_columns
+        ] + real_output_columns
 
     valid_dataset_ids = []
     for dataset in datasets:

DashAI/back/api/api_v1/endpoints/model_sessions.py

@@ -14,6 +14,7 @@
     Prediction,
     Run,
 )
+from DashAI.back.job.base_job import JobError
 from DashAI.back.job.predict_job import run_manual_prediction
 
 if TYPE_CHECKING:
@@ -154,6 +155,9 @@ async def filter_datasets_endpoint(
         ``{"valid_dataset_ids": [...]}`` with the ids of the matching datasets.
     """
     from DashAI.back.dataloaders.classes.dashai_dataset import get_columns_spec
+    from DashAI.back.job.session_preprocessing_job import (
+        get_real_input_output_columns,
+    )
 
     try:
         with session_factory() as db:
@@ -188,7 +192,10 @@ async def filter_datasets_endpoint(
             trained_columns_spec = get_columns_spec(
                 f"{trained_dataset.file_path}/dataset"
             )
-            output_columns = set(model_session.output_columns)
+            _real_input_columns, real_output_columns = get_real_input_output_columns(
+                model_session
+            )
+            output_columns = set(real_output_columns)
             input_columns = [
                 col for col in trained_columns_spec if col not in output_columns
             ]
@@ -209,6 +216,14 @@ async def filter_datasets_endpoint(
     except HTTPException:
         # Re-raise HTTPExceptions as-is
         raise
+    except JobError as e:
+        # A session that isn't fully configured yet (e.g. no real output
+        # column resolved) — a client-actionable state, not an unexpected
+        # internal error.
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        ) from e
     except Exception as e:
         logger.exception("Error filtering datasets: %s", str(e))
         raise HTTPException(

DashAI/back/api/api_v1/endpoints/predict.py

@@ -1,14 +1,52 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Literal, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
+
+
+class ColumnAtom(BaseModel):
+    """A selectable unit in the session wizard: either a literal column of
+    the raw dataset, or the whole output group (one declared slot) of a
+    converter already configured earlier in the same converters list.
+
+    A `group` atom always selects the entire slot: there is no way to pick
+    a subset of a slot's real columns, since their exact names/count are
+    not known until the real fit runs (see the design spec, section B).
+    """
+
+    kind: Literal["column", "group"]
+    name: Optional[str] = None
+    converter_id: Optional[str] = None
+    slot: Optional[int] = None
+
+
+def _coerce_atom(value: Any) -> Any:
+    """Coerce a bare column-name string into a literal `column` atom.
+
+    Kept for backward compatibility: callers that predate the atom-based
+    column scope (e.g. `input_columns=["SepalLengthCm"]`) still POST plain
+    strings. Anything that isn't a string (a dict, or an already-built
+    `ColumnAtom`) is passed through unchanged for pydantic to validate
+    normally.
+    """
+    if isinstance(value, str):
+        return ColumnAtom(kind="column", name=value)
+    return value
 
 
 class SessionConverterParams(BaseModel):
+    id: str
     converter: str
     params: Dict[str, Any] = {}
-    columns: List[str] = []
+    input_scope: List[ColumnAtom] = []
     target_column: Optional[str] = None
 
+    @field_validator("input_scope", mode="before")
+    @classmethod
+    def _coerce_input_scope(cls, value: Any) -> Any:
+        if isinstance(value, list):
+            return [_coerce_atom(item) for item in value]
+        return value
+
 
 class UpdateConvertersParams(BaseModel):
     converters: List[SessionConverterParams] = []
@@ -18,15 +56,22 @@ class ModelSessionParams(BaseModel):
     dataset_id: int
     task_name: str
     name: str
-    input_columns: List[str] = []
-    output_columns: List[str] = []
+    input_columns: List[ColumnAtom] = []
+    output_columns: List[ColumnAtom] = []
     train_metrics: List[str]
     validation_metrics: List[str]
     test_metrics: List[str]
     evaluation_strategy: str
     splits: str
     converters: List[SessionConverterParams] = []
 
+    @field_validator("input_columns", "output_columns", mode="before")
+    @classmethod
+    def _coerce_columns(cls, value: Any) -> Any:
+        if isinstance(value, list):
+            return [_coerce_atom(item) for item in value]
+        return value
+
 
 class ColumnsValidationParams(BaseModel):
     task_name: str

DashAI/back/api/api_v1/schemas/model_sessions_params.py

@@ -631,6 +631,24 @@ def select_columns(self, column_names: Union[str, List[str]]) -> "DashAIDataset"
             col: self._types[col] for col in column_names if col in self._types
         }
 
+        # `pa.Table.select()` only prunes columns; it does not touch the
+        # table's own schema-level `dashai_types` metadata blob, which
+        # would otherwise still list every column from the *original*
+        # table, not just the ones kept here. Left uncorrected, saving
+        # this table to disk and reloading it (`save_dataset`/
+        # `load_dataset`, which derive `.types` from that same blob when
+        # no explicit `types=` is given) would silently reintroduce
+        # phantom type entries for columns that no longer exist in the
+        # actual data — crashing the first thing that iterates `.types`
+        # expecting it to match the real columns (e.g.
+        # `categorical_label_encoder`).
+        from DashAI.back.types.utils import save_types_in_arrow_metadata
+
+        subset_table = save_types_in_arrow_metadata(
+            subset_table,
+            {col: t.to_string() for col, t in subset_types.items()},
+        )
+
         return DashAIDataset(table=subset_table, splits=self.splits, types=subset_types)
 
     def __getitem__(self, key):