diff --git a/plexe/workflow.py b/plexe/workflow.py index d81ab93c..13de000b 100644 --- a/plexe/workflow.py +++ b/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) diff --git a/tests/CODE_INDEX.md b/tests/CODE_INDEX.md index 1fbe1b68..a9f9c80c 100644 --- a/tests/CODE_INDEX.md +++ b/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` diff --git a/tests/unit/workflow/test_prepare_data_explicit_splits.py b/tests/unit/workflow/test_prepare_data_explicit_splits.py index 2047dc6c..2026a59f 100644 --- a/tests/unit/workflow/test_prepare_data_explicit_splits.py +++ b/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, + )