feat: add column exclusion pipeline (#181)

* feat: add column exclusion pipeline

* fix: address review feedback

* chore: bump version to 1.3.4

Author: marcellodebernardi

plexe/CODE_INDEX.md

@@ -1,6 +1,6 @@
 # Code Index: plexe
 
-> Generated on 2026-02-27 15:11:29
+> Generated on 2026-03-02 12:38:31
 
 Code structure and public interface documentation for the **plexe** package.
 

plexe/agents/ml_task_analyser.py

@@ -79,7 +79,12 @@ def _build_agent(self) -> CodeAgent:
                 "5. **Data Challenges**: Issues specific to this data type (e.g., path validity, text length variance, missing values)\n"
                 "6. **Preprocessing Recommendations**: High-level guidance on preparing data for modeling\n"
                 "7. **Feature Engineering** (if applicable): Suggested transformations, interactions, encodings\n"
-                "8. **Split Strategy** (in recommended_split):\n"
+                "8. **Problematic Columns**: Identify columns to exclude and categorize them as one of: leakage, constant, identifier, irrelevant.\n"
+                "   - leakage: correlation > 0.95 with target or post-hoc derivatives of target\n"
+                "   - constant: near-zero variance or single unique value\n"
+                "   - identifier: IDs, keys, hashes, or unique identifiers\n"
+                "   - irrelevant: columns unrelated to the task or likely to add noise\n"
+                "9. **Split Strategy** (in recommended_split):\n"
                 "   - temporal_reasoning: Chronological split only if predicting FUTURE time periods. Timestamp as metadata → random OK.\n"
                 "   - stratification_reasoning: Stratify if classification with class imbalance.\n"
                 "\n"

plexe/constants.py

@@ -21,6 +21,7 @@ class ScratchKeys:
     SAVED_MODEL = "_saved_model"
     STATISTICAL_PROFILE = "_statistical_profile"
     EDA_REPORT = "_eda_report"
+    PROBLEMATIC_COLUMNS = "_problematic_columns"
 
     # Dataset URIs
     TRAIN_URI = "_train_uri"

plexe/models.py

@@ -86,6 +86,7 @@ class BuildContext:
     compute_metric: Any | None = None  # Function for computing metric (callable)
     output_targets: list[str] = field(default_factory=list)  # Target column(s) identified by MLTaskAnalyser
     group_column: str | None = None  # For ranking: query_id, session_id, user_id (group identifier)
+    excluded_columns: list[dict] = field(default_factory=list)  # Columns removed before downstream phases
 
     # Data preparation phase
     train_uri: str | None = None
@@ -156,6 +157,7 @@ def to_dict(self) -> dict:
             ),
             "output_targets": self.output_targets,
             "group_column": self.group_column,
+            "excluded_columns": self.excluded_columns,
             # Phase 2 fields
             "train_uri": self.train_uri,
             "val_uri": self.val_uri,
@@ -207,6 +209,7 @@ def from_dict(d: dict) -> "BuildContext":
             compute_metric=compute_metric,
             output_targets=d.get("output_targets", []),
             group_column=d.get("group_column"),
+            excluded_columns=d.get("excluded_columns", []),
             train_uri=d.get("train_uri"),
             val_uri=d.get("val_uri"),
             test_uri=d.get("test_uri"),

plexe/tools/submission.py

@@ -13,7 +13,7 @@
 from sklearn.pipeline import Pipeline
 from smolagents import tool
 
-from plexe.constants import DirNames
+from plexe.constants import DirNames, ScratchKeys
 from plexe.models import BuildContext, Metric, Hypothesis, UnifiedPlan
 from plexe.search.insight_store import InsightStore
 from plexe.utils.tracing import tool_span
@@ -668,6 +668,7 @@ def save_eda_report(
         key_insights: list[str],
         recommended_split: dict[str, Any],
         feature_relationships: dict[str, Any] | None = None,
+        problematic_columns: list[dict] | None = None,
         group_column: str | None = None,
     ) -> str:
         """
@@ -685,6 +686,7 @@ def save_eda_report(
             key_insights: Important ML findings relevant to the task
             recommended_split: Split strategy dict with: ratios (dict), temporal_reasoning (str explaining if/why chronological split needed), stratification_reasoning (str explaining if/why stratified split needed)
             feature_relationships: Optional dict with feature-target relationships (only when computable, e.g., correlations for tabular data)
+            problematic_columns: Optional list of dicts with {"column": str, "reason": str, "category": str} where category in {"leakage", "constant", "identifier", "irrelevant"}
             group_column: Optional group/query ID column for ranking tasks (e.g., "session_id", "query_id", "user_id")
 
         Example (Tabular):
@@ -697,7 +699,8 @@ def save_eda_report(
                 preprocessing_recommendations=["Impute missing values", "Encode categorical features"],
                 key_insights=["CryoSleep is highly predictive based on correlation"],
                 recommended_split={"ratios": {"train": 0.7, "val": 0.15, "test": 0.15}, "temporal_reasoning": "No chronological split needed - this is cross-sectional classification of records, not forecasting future events", "stratification_reasoning": "Stratified split recommended due to class imbalance to maintain balance across splits"},
-                feature_relationships={"correlations_with_target": {"CryoSleep": 0.42, "Age": -0.15}}
+                feature_relationships={"correlations_with_target": {"CryoSleep": 0.42, "Age": -0.15}},
+                problematic_columns=[{"column": "leak_col", "reason": "Correlation > 0.95 with target", "category": "leakage"}]
             )
 
         Example (Image):
@@ -744,6 +747,9 @@ def save_eda_report(
             Confirmation message
         """
 
+        if problematic_columns is None:
+            problematic_columns = []
+
         # Build structured report
         report = {
             "task_type": task_type,
@@ -755,11 +761,13 @@ def save_eda_report(
             "key_insights": key_insights,
             "recommended_split": recommended_split,
             "feature_relationships": feature_relationships,
+            "problematic_columns": problematic_columns,
             "group_column": group_column,
         }
 
         # Save to context
         context.scratch["_eda_report"] = report
+        context.scratch[ScratchKeys.PROBLEMATIC_COLUMNS] = problematic_columns
 
         # Also save group_column directly to context for ranking tasks
         if group_column:

plexe/workflow.py

@@ -35,7 +35,7 @@
 
 from plexe.integrations.base import WorkflowIntegration
 from plexe.config import Config
-from plexe.constants import DirNames, PhaseNames
+from plexe.constants import DirNames, PhaseNames, ScratchKeys
 from plexe.models import BuildContext, Solution, Baseline, Hypothesis, DataLayout, EvaluationReport
 from plexe.execution.training.runner import TrainingRunner
 from plexe.checkpointing import save_checkpoint, load_checkpoint
@@ -204,10 +204,13 @@ def build_model(
         # Phase 2: Data Preparation
         if start_phase <= 2:
             with tracer.start_as_current_span("Phase 2: Data Preparation"):
-                # Use sanitized dataset if available (from Phase 1 column name cleaning)
-                # If resuming from checkpoint and sanitized URI not in context, use original
-                # (Phase 1 checkpoint should have stored the sanitized URI)
-                train_uri_to_use = context.scratch.get("_sanitized_dataset_uri", train_dataset_uri)
+                # Use filtered dataset if available. Note: filtering is applied on top of the
+                # sanitized dataset, so choosing filtered implicitly includes any sanitization.
+                # If no exclusions were applied, filtered == sanitized. If resuming from
+                # checkpoint and URIs not in context, fall back to original.
+                train_uri_to_use = context.scratch.get("_filtered_dataset_uri") or context.scratch.get(
+                    "_sanitized_dataset_uri", train_dataset_uri
+                )
 
                 # Sanitize test dataset too if provided (must match training data schema)
                 test_uri_to_use = test_dataset_uri
@@ -574,6 +577,135 @@ def sanitize_dataset_column_names(spark: SparkSession, dataset_uri: str, context
     return sanitized_uri
 
 
+def _set_noop_filtered_dataset_uri(context: BuildContext, dataset_uri: str) -> str:
+    """Store no-op filtering result in context and return original URI."""
+    context.excluded_columns = []
+    context.scratch["_filtered_dataset_uri"] = dataset_uri
+    return dataset_uri
+
+
+def _get_problematic_columns_payload(context: BuildContext) -> list[object] | None:
+    """Read and validate problematic columns payload from scratch."""
+    problematic_columns = context.scratch.get(ScratchKeys.PROBLEMATIC_COLUMNS, [])
+    if not problematic_columns:
+        logger.info("No problematic columns flagged - skipping exclusion step")
+        return None
+    if not isinstance(problematic_columns, list):
+        logger.warning("Problematic columns payload is not a list - skipping exclusion step")
+        return None
+    return problematic_columns
+
+
+def _build_protected_columns_set(context: BuildContext) -> set[str]:
+    """Build set of columns that can never be excluded."""
+    protected_columns = set(context.output_targets or [])
+    if context.group_column:
+        protected_columns.add(context.group_column)
+    if context.primary_input_column:
+        protected_columns.add(context.primary_input_column)
+    return protected_columns
+
+
+def _normalize_exclusion_reason(raw_reason: object) -> str:
+    """Normalize exclusion reason to a non-empty string."""
+    if isinstance(raw_reason, str) and raw_reason.strip():
+        return raw_reason
+    if not raw_reason:
+        return "unspecified"
+    return str(raw_reason)
+
+
+def _filter_valid_exclusions(
+    problematic_columns: list[object], available_columns: set[str], protected_columns: set[str]
+) -> tuple[list[str], list[dict], list[str], list[str], int]:
+    """Validate exclusion entries and return drop candidates + diagnostics."""
+    columns_to_drop: list[str] = []
+    excluded_entries: list[dict] = []
+    skipped_protected: list[str] = []
+    skipped_missing: list[str] = []
+    invalid_entries = 0
+    seen: set[str] = set()
+
+    for entry in problematic_columns:
+        if not isinstance(entry, dict):
+            invalid_entries += 1
+            continue
+        column = entry.get("column")
+        if not isinstance(column, str) or not column.strip():
+            invalid_entries += 1
+            continue
+        if column in seen:
+            continue
+        seen.add(column)
+        if column in protected_columns:
+            skipped_protected.append(column)
+            continue
+        if column not in available_columns:
+            skipped_missing.append(column)
+            continue
+        columns_to_drop.append(column)
+        excluded_entries.append({"column": column, "reason": _normalize_exclusion_reason(entry.get("reason"))})
+
+    return columns_to_drop, excluded_entries, skipped_protected, skipped_missing, invalid_entries
+
+
+def _exclude_problematic_columns(
+    spark: SparkSession,
+    dataset_uri: str,
+    context: BuildContext,
+    config: Config | None,
+) -> str:
+    """
+    Drop problematic columns identified during Phase 1 analysis.
+
+    Args:
+        spark: SparkSession
+        dataset_uri: Dataset URI to filter (already sanitized)
+        context: Build context with problematic columns and targets
+        config: Configuration (reserved for future use)
+
+    Returns:
+        URI of filtered dataset (or original if no exclusions needed)
+    """
+    _ = config
+    problematic_columns = _get_problematic_columns_payload(context)
+    if problematic_columns is None:
+        return _set_noop_filtered_dataset_uri(context, dataset_uri)
+
+    df = spark.read.parquet(dataset_uri)
+    columns_to_drop, excluded_entries, skipped_protected, skipped_missing, invalid_entries = _filter_valid_exclusions(
+        problematic_columns=problematic_columns,
+        available_columns=set(df.columns),
+        protected_columns=_build_protected_columns_set(context),
+    )
+
+    if skipped_protected:
+        logger.warning(
+            "Problematic columns include protected columns; skipping exclusions for: "
+            + ", ".join(sorted(skipped_protected))
+        )
+    if skipped_missing:
+        logger.warning(
+            "Problematic columns not found in dataset; skipping exclusions for: " + ", ".join(sorted(skipped_missing))
+        )
+    if invalid_entries:
+        logger.warning(f"Skipped {invalid_entries} malformed problematic column entries")
+    if not columns_to_drop:
+        logger.info("No valid problematic columns to exclude after validation")
+        return _set_noop_filtered_dataset_uri(context, dataset_uri)
+
+    logger.info(f"Excluding {len(columns_to_drop)} problematic columns from dataset")
+    for entry in excluded_entries:
+        logger.info(f"  drop '{entry['column']}': {entry['reason']}")
+
+    filtered_uri = f"{context.work_dir}/{DirNames.BUILD_DIR}/data/dataset_filtered.parquet"
+    df.drop(*columns_to_drop).write.mode("overwrite").parquet(filtered_uri)
+    logger.info(f"✓ Filtered dataset saved: {filtered_uri}")
+    context.excluded_columns = excluded_entries
+    context.scratch["_filtered_dataset_uri"] = filtered_uri
+    return filtered_uri
+
+
 def analyze_data(
     spark: SparkSession,
     dataset_uri: str,
@@ -642,6 +774,9 @@ def analyze_data(
     save_report(context.work_dir, "02_task_analysis", task_analysis)
     # Note: output_targets already set by task_agent.run()
 
+    # Step 3b: Exclude problematic columns identified during analysis
+    _exclude_problematic_columns(spark, sanitized_uri, context, config)
+
     # Step 3: Metric Selection
     metric_agent = MetricSelectorAgent(context, config)
     metric = metric_agent.run()
@@ -711,6 +846,16 @@ def prepare_data(
 
         # Copy test dataset to DirNames.BUILD_DIR/data/ for consistency
         test_df = spark.read.parquet(test_dataset_uri)
+        if context.excluded_columns:
+            excluded_column_names = [
+                entry.get("column")
+                for entry in context.excluded_columns
+                if isinstance(entry, dict) and entry.get("column")
+            ]
+            columns_to_drop = [col for col in excluded_column_names if col in test_df.columns]
+            if columns_to_drop:
+                test_df = test_df.drop(*columns_to_drop)
+                logger.info(f"Excluded {len(columns_to_drop)} columns from test dataset: {columns_to_drop}")
         test_uri = str(context.work_dir / DirNames.BUILD_DIR / "data" / "test.parquet")
         test_df.write.mode("overwrite").parquet(test_uri)
         logger.info(f"Copied test dataset to: {test_uri}")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "plexe"
-version = "1.3.3"
+version = "1.3.4"
 description = "An agentic framework for building ML models from natural language"
 authors = [
     "Marcello De Bernardi <mdebernardi@plexe.ai>",

tests/CODE_INDEX.md

@@ -1,6 +1,6 @@
 # Code Index: tests
 
-> Generated on 2026-02-27 15:11:29
+> Generated on 2026-03-02 12:38:31
 
 Test suite structure and test case documentation.
 
@@ -169,6 +169,18 @@ Unit tests for validation functions.
 - `test_validate_metric_function_object_success()` - Callable with correct signature should pass.
 - `test_validate_metric_function_object_bad_signature()` - Callable with wrong arg names should fail.
 
+---
+## `unit/workflow/test_column_exclusion.py`
+Tests for column exclusion pipeline.
+
+**Functions:**
+- `spark()` - No description
+- `test_exclude_problematic_columns_drops_columns_and_returns_new_uri(spark, tmp_path)` - No description
+- `test_exclude_problematic_columns_noop_when_empty(spark, tmp_path)` - No description
+- `test_build_context_round_trip_with_excluded_columns(tmp_path)` - No description
+- `test_exclude_problematic_columns_never_drops_target(spark, tmp_path)` - No description
+- `test_exclude_problematic_columns_never_drops_primary_input(spark, tmp_path)` - No description
+
 ---
 ## `unit/workflow/test_model_card.py`
 Unit tests for model card generation.