fix(data): validate explicit split columns

plexe/workflow.py

@@ -911,6 +911,22 @@ def _materialize_explicit_split(
     return output_uri
 
 
+def _validate_explicit_split_columns(
+    spark: SparkSession, training_dataset_uri: str, split_uris: dict[str, str | None]
+) -> None:
+    """Ensure each explicit split contains every column present in training data."""
+    explicit_splits = {name: uri for name, uri in split_uris.items() if uri is not None}
+    if not explicit_splits:
+        return
+
+    training_columns = set(spark.read.parquet(training_dataset_uri).columns)
+    for split_name, split_uri in explicit_splits.items():
+        split_columns = set(spark.read.parquet(split_uri).columns)
+        missing_columns = sorted(training_columns - split_columns)
+        if missing_columns:
+            raise ValueError(f"Explicit {split_name} split is missing training columns: {', '.join(missing_columns)}")
+
+
 def prepare_data(
     spark: SparkSession,
     training_dataset_uri: str,
@@ -957,6 +973,12 @@ def prepare_data(
     )
     train_uri = training_dataset_uri
 
+    _validate_explicit_split_columns(
+        spark,
+        training_dataset_uri,
+        {"validation": val_uri if provided_val else None, "test": test_uri if provided_test else None},
+    )
+
     def run_split(split_ratios: dict[str, float], output_dir: str = splits_output_dir) -> tuple[str, str, str | None]:
         splitter = DatasetSplitterAgent(spark=spark, dataset_uri=training_dataset_uri, context=context, config=config)
         return splitter.run(split_ratios=split_ratios, output_dir=output_dir)

tests/CODE_INDEX.md

@@ -1,6 +1,6 @@
 # Code Index: tests
 
-> Generated on 2026-03-05 21:32:55
+> Generated on 2026-10-10 05:18:37
 
 Test suite structure and test case documentation.
 
@@ -448,6 +448,7 @@ Unit tests for prepare_data split resolution with explicit val/test datasets.
 - `test_prepare_data_uses_all_provided_splits_without_running_splitter(monkeypatch, tmp_path)` - No description
 - `test_prepare_data_generates_missing_test_when_only_val_is_provided(monkeypatch, tmp_path)` - No description
 - `test_prepare_data_generates_missing_val_when_only_test_is_provided(monkeypatch, tmp_path)` - No description
+- `test_prepare_data_rejects_explicit_split_missing_training_column(monkeypatch, tmp_path, split_name, val_dataset_uri, test_dataset_uri, materialized_uri, missing_column)` - No description
 
 ---
 ## `unit/workflow/test_resume_model_type_filtering.py`

tests/unit/workflow/test_prepare_data_explicit_splits.py

@@ -7,7 +7,6 @@
 import pytest
 
 pytest.importorskip("dataclasses_json")
-pytest.importorskip("pyspark")
 
 from plexe.models import BuildContext
 import plexe.workflow as workflow
@@ -23,6 +22,24 @@ def ensure_local(self, uris, work_dir):  # noqa: D401
         return uris
 
 
+class _SchemaFrame:
+    def __init__(self, columns: list[str]):
+        self.columns = columns
+
+
+class _SchemaReader:
+    def __init__(self, schemas: dict[str, list[str]]):
+        self._schemas = schemas
+
+    def parquet(self, uri: str) -> _SchemaFrame:
+        return _SchemaFrame(self._schemas[uri])
+
+
+class _SchemaSpark:
+    def __init__(self, schemas: dict[str, list[str]]):
+        self.read = _SchemaReader(schemas)
+
+
 def _make_context(tmp_path) -> BuildContext:
     context = BuildContext(
         user_id="user-1",
@@ -65,7 +82,13 @@ def run(self, train_uri, val_uri, train_sample_size, val_sample_size, output_dir
     monkeypatch.setattr(workflow, "_save_phase_checkpoint", lambda *args, **kwargs: None)
 
     workflow.prepare_data(
-        spark=object(),
+        spark=_SchemaSpark(
+            {
+                "train_input.parquet": ["feature", "target"],
+                "copied_val.parquet": ["feature", "target", "val_only_column"],
+                "copied_test.parquet": ["feature", "target", "test_only_column"],
+            }
+        ),
         training_dataset_uri="train_input.parquet",
         val_dataset_uri="val_input.parquet",
         test_dataset_uri="test_input.parquet",
@@ -116,7 +139,7 @@ def run(self, train_uri, val_uri, train_sample_size, val_sample_size, output_dir
     monkeypatch.setattr(workflow, "_save_phase_checkpoint", lambda *args, **kwargs: None)
 
     workflow.prepare_data(
-        spark=object(),
+        spark=_SchemaSpark({"train_input.parquet": ["feature", "target"], "copied_val.parquet": ["feature", "target"]}),
         training_dataset_uri="train_input.parquet",
         val_dataset_uri="val_input.parquet",
         test_dataset_uri=None,
@@ -165,7 +188,9 @@ def run(self, train_uri, val_uri, train_sample_size, val_sample_size, output_dir
     monkeypatch.setattr(workflow, "_save_phase_checkpoint", lambda *args, **kwargs: None)
 
     workflow.prepare_data(
-        spark=object(),
+        spark=_SchemaSpark(
+            {"train_input.parquet": ["feature", "target"], "copied_test.parquet": ["feature", "target"]}
+        ),
         training_dataset_uri="train_input.parquet",
         val_dataset_uri=None,
         test_dataset_uri="test_input.parquet",
@@ -179,3 +204,49 @@ def run(self, train_uri, val_uri, train_sample_size, val_sample_size, output_dir
     assert context.train_uri == "split_train.parquet"
     assert context.val_uri == "split_val.parquet"
     assert context.test_uri == "copied_test.parquet"
+
+
+@pytest.mark.parametrize(
+    ("split_name", "val_dataset_uri", "test_dataset_uri", "materialized_uri", "missing_column"),
+    [
+        ("validation", "val_input.parquet", None, "copied_val.parquet", "target"),
+        ("validation", "val_input.parquet", None, "copied_val.parquet", "feature"),
+        ("test", None, "test_input.parquet", "copied_test.parquet", "target"),
+        ("test", None, "test_input.parquet", "copied_test.parquet", "feature"),
+    ],
+)
+def test_prepare_data_rejects_explicit_split_missing_training_column(
+    monkeypatch, tmp_path, split_name, val_dataset_uri, test_dataset_uri, materialized_uri, missing_column
+):
+    context = _make_context(tmp_path)
+    config = SimpleNamespace(train_sample_size=100, val_sample_size=40)
+    schemas = {
+        "train_input.parquet": ["feature", "target"],
+        materialized_uri: [column for column in ("feature", "target") if column != missing_column],
+    }
+
+    monkeypatch.setattr(
+        workflow,
+        "_materialize_explicit_split",
+        lambda _spark, _uri, _name, _context, _output_dir: materialized_uri,
+    )
+
+    class _UnexpectedDownstreamWork:
+        def __init__(self, *args, **kwargs):
+            raise AssertionError("Explicit split schemas must be validated before downstream work")
+
+    monkeypatch.setattr(workflow, "DatasetSplitterAgent", _UnexpectedDownstreamWork)
+    monkeypatch.setattr(workflow, "SamplingAgent", _UnexpectedDownstreamWork)
+    monkeypatch.setattr(workflow, "_save_phase_checkpoint", lambda *args, **kwargs: None)
+
+    with pytest.raises(ValueError, match=f"{split_name}.*{missing_column}"):
+        workflow.prepare_data(
+            spark=_SchemaSpark(schemas),
+            training_dataset_uri="train_input.parquet",
+            val_dataset_uri=val_dataset_uri,
+            test_dataset_uri=test_dataset_uri,
+            context=context,
+            config=config,
+            integration=_DummyIntegration(),
+            generate_test_set=True,
+        )