Feat session preprocessing structure

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

@@ -11,11 +11,17 @@
     ColumnsValidationParams,
     ModelSessionBulkDeleteParams,
     ModelSessionParams,
+    PreprocessingStructureParams,
 )
 from DashAI.back.api.utils import remove_path
 from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run
 from DashAI.back.job.preprocessing_job import PreprocessingJob
 from DashAI.back.preprocessing.column_ref import ConverterSequence, RawColumnRef
+from DashAI.back.preprocessing.structure import (
+    StructureError,
+    infer_structure,
+    resolve_state_refs,
+)
 from DashAI.back.splitters.splits_payload import (
     META_KEYS,
     normalize_splits_payload,
@@ -121,40 +127,119 @@ async def get_model_session(
         return model_session
 
 
-@router.post("/validation")
-@inject
-async def validate_columns(
-    params: ColumnsValidationParams,
-    component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]),
-    session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]),
-):
-    """Validate if dataset columns are compatible with a task."""
+def _load_sample_dataset(db, dataset_id: int):
+    """Load the first rows of a dataset, enough for its columns and types.
+
+    Parameters
+    ----------
+    db : Session
+        An open database session.
+    dataset_id : int
+        The dataset to load.
+
+    Returns
+    -------
+    DashAIDataset
+        At most 5 rows of the dataset, with every column and its type.
+
+    Raises
+    ------
+    HTTPException
+        404 if the dataset is not registered.
+    """
     import os
 
     import pyarrow as pa
     import pyarrow.ipc as ipc
 
     from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 
+    dataset = db.get(Dataset, dataset_id)
+    if not dataset:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Dataset not found",
+        )
+    data_filepath = os.path.join(f"{dataset.file_path}/dataset", "data.arrow")
+    with pa.OSFile(data_filepath, "rb") as source:
+        reader = ipc.open_file(source)
+        batch = reader.get_batch(0)
+        sample_batch = batch.slice(0, min(5, batch.num_rows))
+    return DashAIDataset(pa.Table.from_batches([sample_batch]))
+
+
+def _type_allowed(type_name, allowed_types) -> bool:
+    """Whether a column type is one of a task's allowed input types.
+
+    A task declaring the DashAIValue wildcard accepts every concrete value
+    type; an unknown type (None) is let through, since only a fit can tell.
+    """
+    if not allowed_types or type_name is None:
+        return True
+    return type_name in allowed_types or (
+        "DashAIValue" in allowed_types and type_name in _DASHAI_VALUE_TYPE_NAMES
+    )
+
+
+@router.post("/preprocessing/structure")
+@inject
+async def get_preprocessing_structure(
+    params: PreprocessingStructureParams,
+    component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]),
+    session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]),
+):
+    """Estimate the dataset structure after every step of a converter chain.
+
+    Read only: nothing is fit and nothing is stored. The session wizard calls
+    it on every change to the chain, to offer only columns that exist at each
+    step and to flag steps that cannot work.
+
+    Parameters
+    ----------
+    params : PreprocessingStructureParams
+        The dataset, the candidate and target columns, and the chain.
+
+    Returns
+    -------
+    StructureResult
+        The state after every step, their errors and warnings, and the final
+        state.
+    """
     with session_factory() as db:
         try:
-            dataset = db.get(Dataset, params.dataset_id)
-            if not dataset:
-                raise HTTPException(
-                    status_code=status.HTTP_404_NOT_FOUND,
-                    detail="Dataset not found",
-                )
+            sample = _load_sample_dataset(db, params.dataset_id)
+        except exc.SQLAlchemyError as e:
+            log.exception(e)
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal database error",
+            ) from e
+    return infer_structure(
+        dataset_types=sample.types,
+        candidates=params.candidates,
+        target=params.output_columns,
+        steps=params.steps,
+        component_registry=component_registry,
+    )
 
-            dataset_path = f"{dataset.file_path}/dataset"
-            data_filepath = os.path.join(dataset_path, "data.arrow")
-            with pa.OSFile(data_filepath, "rb") as source:
-                reader = ipc.open_file(source)
-                batch = reader.get_batch(0)
-                sample_size = min(5, batch.num_rows)
-                sample_batch = batch.slice(0, sample_size)
 
-            table = pa.Table.from_batches([sample_batch])
-            minimal_dataset = DashAIDataset(table)
+@router.post("/validation")
+@inject
+async def validate_columns(
+    params: ColumnsValidationParams,
+    component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]),
+    session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]),
+):
+    """Validate if dataset columns are compatible with a task.
+
+    With `preprocessing`, every input ref is resolved against the chain's
+    estimated final structure and its type checked against the task's input
+    types; the original columns among them also go through the task's own
+    validation, as without preprocessing.
+    """
+    with session_factory() as db:
+        try:
+            minimal_dataset = _load_sample_dataset(db, params.dataset_id)
 
             column_names = minimal_dataset.column_names
 
@@ -192,31 +277,42 @@ async def validate_columns(
 
     task: "BaseTask" = component_registry[params.task_name]["class"]()
 
-    if group_refs:
-        declared_types = params.converter_output_types or {}
-        task_metadata = task.get_metadata()
-        allowed_input_types = set(task_metadata.get("inputs_types", []))
-        for ref in group_refs:
-            # A step with a heterogeneous scope (see SessionPreprocessor.
-            # _classify_by_type) can declare more than one type, one per
-            # slot — "{step}:{slot}" disambiguates which one a slotted ref
-            # means; an unslotted ref (the whole step) keeps the plain
-            # "{step}" key, unchanged from before slots existed.
-            key = str(ref.step) if ref.slot is None else f"{ref.step}:{ref.slot}"
-            declared_type = declared_types.get(key)
-            type_ok = declared_type in allowed_input_types or (
-                "DashAIValue" in allowed_input_types
-                and declared_type in _DASHAI_VALUE_TYPE_NAMES
+    if params.preprocessing:
+        allowed_input_types = set(task.get_metadata().get("inputs_types", []))
+        structure = infer_structure(
+            dataset_types=minimal_dataset.types,
+            candidates=list(column_names),
+            target=outputs_names,
+            steps=params.preprocessing,
+            component_registry=component_registry,
+        )
+        if not structure.valid:
+            return {
+                "dataset_status": "invalid",
+                "error": "The preprocessing chain has an invalid step.",
+            }
+        try:
+            items = resolve_state_refs(
+                params.input_refs or [], structure.final, set(outputs_names)
             )
-            if allowed_input_types and not type_ok:
+        except StructureError as e:
+            return {"dataset_status": "invalid", "error": e.message.code}
+        for item in items:
+            if not _type_allowed(item.type, allowed_input_types):
                 return {
                     "dataset_status": "invalid",
                     "error": (
-                        f"Converter step {ref.step} declares output type "
-                        f"'{declared_type}', which is not one of the task's "
+                        f"Input of type '{item.type}' is not one of the task's "
                         f"allowed input types {sorted(allowed_input_types)}."
                     ),
                 }
+    elif group_refs:
+        # A ref to a converter's output only has a type through the chain's
+        # estimated structure, which needs the preprocessing steps.
+        return {
+            "dataset_status": "invalid",
+            "error": "References to converter outputs need the preprocessing steps.",
+        }
 
     validation_response = {}
 
@@ -286,6 +382,52 @@ def _validate_splits(splits: str, component_registry: "ComponentRegistry") -> No
         ) from e
 
 
+def _check_preprocessing_structure(
+    db, params, sequence, input_column_refs, component_registry
+) -> None:
+    """Reject a session whose chain or inputs cannot work, before creating it.
+
+    Without this, the session would be created and its PreprocessingJob
+    would only fail afterwards, with an error like a bare missing column
+    name.
+
+    Raises
+    ------
+    HTTPException
+        422 with the estimated structure when a step is invalid, or with the
+        error when an input ref points at something the chain does not
+        produce (or consumed).
+    """
+    sample = _load_sample_dataset(db, params.dataset_id)
+    structure = infer_structure(
+        dataset_types=sample.types,
+        candidates=list(sample.column_names),
+        target=params.output_columns,
+        steps=sequence.steps,
+        component_registry=component_registry,
+    )
+    if not structure.valid:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail={
+                "message": "The preprocessing chain has an invalid step.",
+                "structure": structure.model_dump(mode="json"),
+            },
+        )
+    try:
+        resolve_state_refs(
+            input_column_refs, structure.final, set(params.output_columns)
+        )
+    except StructureError as e:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail={
+                "message": "An input column does not exist after preprocessing.",
+                "error": e.message.model_dump(mode="json"),
+            },
+        ) from e
+
+
 @router.post("/", status_code=status.HTTP_201_CREATED)
 @inject
 async def create_model_session(
@@ -362,6 +504,11 @@ async def create_model_session(
                 schema = reader.schema
                 column_names = schema.names
 
+            if has_preprocessing:
+                _check_preprocessing_structure(
+                    db, params, sequence, input_column_refs, component_registry
+                )
+
             # When preprocessing is configured, input_columns is only a
             # placeholder until PreprocessingJob resolves the real ones, so
             # it cannot be checked against the raw dataset's column count.

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

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import List, Optional
 
 from pydantic import BaseModel, Field
 
@@ -26,7 +26,17 @@ class ColumnsValidationParams(BaseModel):
     inputs_columns: List[str]
     outputs_columns: List[str]
     input_refs: Optional[List[ColumnRef]] = None
-    converter_output_types: Optional[Dict[str, str]] = None
+    # With preprocessing, the type of every input ref is taken from the
+    # chain's estimated structure (see infer_structure).
+    preprocessing: Optional[List[ConverterStep]] = None
+
+
+class PreprocessingStructureParams(BaseModel):
+    dataset_id: int
+    # Original columns the user may feed into the chain.
+    candidates: List[str]
+    output_columns: List[str]
+    steps: List[ConverterStep] = Field(default_factory=list)
 
 
 class ModelSessionBulkDeleteParams(BaseModel):