The Schema Change Problem in ML Pipelines
ML pipelines depend on data schemas—the structure, types, and constraints of input data. When a data engineering team adds a column to a source table, changes a field type from string to integer, or renames a feature, the downstream ML pipeline can fail silently (producing wrong predictions) or loudly (crashing the training job). Unlike traditional software where schema changes are managed through database migrations, ML pipelines span multiple systems (data warehouse, feature store, model training, inference API) where a single schema change can cascade through the entire stack.
This playbook provides a systematic approach to managing schema evolution in production ML systems, covering detection, validation, backward compatibility, and automated migration.
Phase 1: Schema Registry
Central Schema Definition
# schema_registry/definitions.py — Central schema definitions for ML pipelines
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any
from enum import Enum
import json
from datetime import datetime
class FieldType(Enum):
INT32 = "int32"
INT64 = "int64"
FLOAT32 = "float32"
FLOAT64 = "float64"
STRING = "string"
BOOLEAN = "boolean"
TIMESTAMP = "timestamp"
ARRAY = "array"
STRUCT = "struct"
class CompatibilityLevel(Enum):
BACKWARD = "backward" # New schema can read old data
FORWARD = "forward" # Old schema can read new data
FULL = "full" # Both backward and forward compatible
NONE = "none" # No compatibility guarantees
@dataclass
class FieldSchema:
name: str
field_type: FieldType
nullable: bool = True
default_value: Any = None
description: str = ""
metadata: Dict[str, str] = field(default_factory=dict)
# For ML-specific metadata
feature_type: Optional[str] = None # "numeric", "categorical", "text", "embedding"
is_target: bool = False
is_key: bool = False
@dataclass
class DatasetSchema:
name: str
version: str
fields: List[FieldSchema]
compatibility: CompatibilityLevel = CompatibilityLevel.BACKWARD
created_at: str = ""
description: str = ""
owner: str = ""
def __post_init__(self):
if not self.created_at:
self.created_at = datetime.utcnow().isoformat()
def get_field(self, name: str) -> Optional[FieldSchema]:
return next((f for f in self.fields if f.name == name), None)
def to_dict(self) -> dict:
return {
"name": self.name,
"version": self.version,
"compatibility": self.compatibility.value,
"fields": [
{
"name": f.name,
"type": f.field_type.value,
"nullable": f.nullable,
"default": f.default_value,
"feature_type": f.feature_type,
"is_target": f.is_target,
"is_key": f.is_key,
}
for f in self.fields
],
}
class SchemaRegistry:
"""
Central registry for all dataset schemas in the ML platform.
Tracks schema versions, validates changes, and enforces compatibility.
"""
def __init__(self, storage_path: str = "schemas/"):
self.storage_path = storage_path
self.schemas: Dict[str, List[DatasetSchema]] = {} # name -> versions
self._load_schemas()
def register_schema(self, schema: DatasetSchema) -> bool:
"""
Register a new schema version. Validates compatibility with previous version.
Returns True if registered, False if compatibility check fails.
"""
name = schema.name
if name in self.schemas and self.schemas[name]:
previous = self.schemas[name][-1]
# Validate compatibility
compatibility_report = self.check_compatibility(previous, schema)
if not compatibility_report["compatible"]:
print(f"❌ Schema change is NOT compatible:")
for issue in compatibility_report["issues"]:
print(f" - {issue}")
return False
print(f"✅ Schema change is compatible ({schema.compatibility.value})")
for change in compatibility_report["changes"]:
print(f" - {change}")
if name not in self.schemas:
self.schemas[name] = []
self.schemas[name].append(schema)
self._save_schema(schema)
return True
def check_compatibility(self, old_schema: DatasetSchema,
new_schema: DatasetSchema) -> dict:
"""Check if new schema is compatible with old schema."""
issues = []
changes = []
old_fields = {f.name: f for f in old_schema.fields}
new_fields = {f.name: f for f in new_schema.fields}
# Check for removed fields
removed = set(old_fields.keys()) - set(new_fields.keys())
for field_name in removed:
old_field = old_fields[field_name]
if old_field.is_target or old_field.is_key:
issues.append(f"Cannot remove target/key field: {field_name}")
elif new_schema.compatibility in [CompatibilityLevel.BACKWARD, CompatibilityLevel.FULL]:
# Backward compatible: removed fields must have been nullable or have defaults
if not old_field.nullable and old_field.default_value is None:
issues.append(
f"Removed non-nullable field without default: {field_name}"
)
else:
changes.append(f"Removed field: {field_name} (safe — was nullable)")
# Check for added fields
added = set(new_fields.keys()) - set(old_fields.keys())
for field_name in added:
new_field = new_fields[field_name]
if new_schema.compatibility in [CompatibilityLevel.FORWARD, CompatibilityLevel.FULL]:
if not new_field.nullable and new_field.default_value is None:
issues.append(
f"Added non-nullable field without default: {field_name}"
)
else:
changes.append(f"Added field: {field_name} (nullable with default)")
else:
changes.append(f"Added field: {field_name}")
# Check for type changes
common = set(old_fields.keys()) & set(new_fields.keys())
for field_name in common:
old_field = old_fields[field_name]
new_field = new_fields[field_name]
if old_field.field_type != new_field.field_type:
if self._is_safe_type_promotion(old_field.field_type, new_field.field_type):
changes.append(
f"Type promotion: {field_name} "
f"{old_field.field_type.value} → {new_field.field_type.value}"
)
else:
issues.append(
f"Unsafe type change: {field_name} "
f"{old_field.field_type.value} → {new_field.field_type.value}"
)
# Check nullable changes
if old_field.nullable and not new_field.nullable:
issues.append(
f"Field {field_name} changed from nullable to non-nullable"
)
return {
"compatible": len(issues) == 0,
"issues": issues,
"changes": changes,
"fields_added": len(added),
"fields_removed": len(removed),
"fields_modified": len([f for f in common
if old_fields[f].field_type != new_fields[f].field_type]),
}
def _is_safe_type_promotion(self, old_type: FieldType, new_type: FieldType) -> bool:
"""Check if a type change is a safe widening promotion."""
safe_promotions = {
(FieldType.INT32, FieldType.INT64),
(FieldType.INT32, FieldType.FLOAT64),
(FieldType.INT64, FieldType.FLOAT64),
(FieldType.FLOAT32, FieldType.FLOAT64),
}
return (old_type, new_type) in safe_promotions
def _load_schemas(self):
"""Load schemas from storage."""
pass # Implementation depends on storage backend
def _save_schema(self, schema: DatasetSchema):
"""Save schema to storage."""
passPhase 2: Schema Validation in Pipelines
Runtime Schema Validation
# schema_validation/validator.py — Validate data against schemas at pipeline boundaries
import pandas as pd
import numpy as np
from typing import Dict, List, Optional
from great_expectations.core import ExpectationSuite
import great_expectations as ge
class MLDataValidator:
"""Validate data against registered schemas at pipeline boundaries."""
def __init__(self, schema_registry: SchemaRegistry):
self.registry = schema_registry
def validate_dataframe(self, df: pd.DataFrame, schema_name: str,
schema_version: str = "latest") -> Dict:
"""
Validate a DataFrame against a registered schema.
Returns validation report with any issues found.
"""
schema = self._get_schema(schema_name, schema_version)
issues = []
warnings = []
# 1. Check all required fields are present
for field in schema.fields:
if field.name not in df.columns:
if not field.nullable and field.default_value is None:
issues.append(f"Missing required field: {field.name}")
else:
warnings.append(f"Missing optional field: {field.name}")
# 2. Check field types
for field in schema.fields:
if field.name not in df.columns:
continue
actual_type = self._pandas_to_field_type(df[field.name].dtype)
if actual_type != field.field_type:
if not self._is_safe_type_promotion(field.field_type, actual_type):
issues.append(
f"Type mismatch for {field.name}: "
f"expected {field.field_type.value}, got {actual_type.value}"
)
# 3. Check null constraints
for field in schema.fields:
if field.name not in df.columns:
continue
if not field.nullable:
null_count = df[field.name].isnull().sum()
if null_count > 0:
issues.append(
f"Non-nullable field {field.name} has {null_count} null values "
f"({null_count/len(df)*100:.1f}%)"
)
# 4. Statistical validation (detect distribution shifts)
for field in schema.fields:
if field.name not in df.columns:
continue
if field.field_type in [FieldType.FLOAT32, FieldType.FLOAT64, FieldType.INT32, FieldType.INT64]:
stats = self._compute_field_stats(df[field.name])
field_stats_key = f"{schema_name}.{field.name}"
# Compare against historical stats
historical = self._get_historical_stats(field_stats_key)
if historical:
drift_score = self._compute_drift_score(stats, historical)
if drift_score > 0.2: # PSI threshold
warnings.append(
f"Possible distribution shift in {field.name}: "
f"PSI={drift_score:.3f}"
)
# 5. Check for unexpected columns
expected_fields = {f.name for f in schema.fields}
unexpected = set(df.columns) - expected_fields
if unexpected:
warnings.append(f"Unexpected columns: {unexpected}")
return {
"valid": len(issues) == 0,
"schema_name": schema_name,
"schema_version": schema.version,
"row_count": len(df),
"issues": issues,
"warnings": warnings,
}
def validate_at_boundary(self, df: pd.DataFrame, schema_name: str,
boundary: str) -> pd.DataFrame:
"""
Validate data at a pipeline boundary (ingestion, feature store, model input).
Raises an exception if validation fails.
"""
report = self.validate_dataframe(df, schema_name)
if not report["valid"]:
error_msg = f"Schema validation failed at {boundary}:\n"
error_msg += "\n".join(f" ❌ {issue}" for issue in report["issues"])
raise SchemaValidationError(error_msg)
if report["warnings"]:
for warning in report["warnings"]:
print(f" ⚠️ {warning}")
return df
def _pandas_to_field_type(self, dtype) -> FieldType:
"""Map pandas dtype to FieldType."""
mapping = {
"int32": FieldType.INT32,
"int64": FieldType.INT64,
"float32": FieldType.FLOAT32,
"float64": FieldType.FLOAT64,
"object": FieldType.STRING,
"bool": FieldType.BOOLEAN,
"datetime64[ns]": FieldType.TIMESTAMP,
}
return mapping.get(str(dtype), FieldType.STRING)
def _compute_field_stats(self, series: pd.Series) -> dict:
return {
"mean": float(series.mean()),
"std": float(series.std()),
"min": float(series.min()),
"max": float(series.max()),
"null_pct": float(series.isnull().mean()),
}
class SchemaValidationError(Exception):
passPhase 3: Automated Schema Migration
Migration Framework
# schema_evolution/migrator.py — Automated schema migration for ML data
from typing import Callable, Dict, List
import pandas as pd
class SchemaMigrator:
"""
Migrate data between schema versions.
Handles field additions, removals, type changes, and renames.
"""
def __init__(self):
self.migrations: Dict[str, List[MigrationStep]] = {}
def register_migration(self, schema_name: str, from_version: str,
to_version: str, steps: List['MigrationStep']):
"""Register a migration path between schema versions."""
key = f"{schema_name}:{from_version}->{to_version}"
self.migrations[key] = steps
def migrate(self, df: pd.DataFrame, schema_name: str,
from_version: str, to_version: str) -> pd.DataFrame:
"""Apply migration steps to transform data from one schema to another."""
key = f"{schema_name}:{from_version}->{to_version}"
if key not in self.migrations:
raise ValueError(f"No migration registered for {key}")
result = df.copy()
for step in self.migrations[key]:
result = step.apply(result)
print(f" Applied: {step.description}")
return result
class MigrationStep:
"""Base class for schema migration steps."""
def __init__(self, description: str):
self.description = description
def apply(self, df: pd.DataFrame) -> pd.DataFrame:
raise NotImplementedError
class AddFieldStep(MigrationStep):
"""Add a new field with a default value or computed expression."""
def __init__(self, field_name: str, default_value=None,
compute_fn: Callable = None, description: str = ""):
super().__init__(description or f"Add field: {field_name}")
self.field_name = field_name
self.default_value = default_value
self.compute_fn = compute_fn
def apply(self, df: pd.DataFrame) -> pd.DataFrame:
if self.compute_fn:
df[self.field_name] = df.apply(self.compute_fn, axis=1)
else:
df[self.field_name] = self.default_value
return df
class RemoveFieldStep(MigrationStep):
"""Remove a field from the schema."""
def __init__(self, field_name: str):
super().__init__(f"Remove field: {field_name}")
self.field_name = field_name
def apply(self, df: pd.DataFrame) -> pd.DataFrame:
if self.field_name in df.columns:
df = df.drop(columns=[self.field_name])
return df
class RenameFieldStep(MigrationStep):
"""Rename a field."""
def __init__(self, old_name: str, new_name: str):
super().__init__(f"Rename field: {old_name} → {new_name}")
self.old_name = old_name
self.new_name = new_name
def apply(self, df: pd.DataFrame) -> pd.DataFrame:
if self.old_name in df.columns:
df = df.rename(columns={self.old_name: self.new_name})
return df
class CastFieldStep(MigrationStep):
"""Change a field's data type."""
def __init__(self, field_name: str, new_type: str):
super().__init__(f"Cast field: {field_name} → {new_type}")
self.field_name = field_name
self.new_type = new_type
def apply(self, df: pd.DataFrame) -> pd.DataFrame:
if self.field_name in df.columns:
df[self.field_name] = df[self.field_name].astype(self.new_type)
return df
# Example: Register migrations
migrator = SchemaMigrator()
migrator.register_migration(
schema_name="customer_features",
from_version="1.0",
to_version="2.0",
steps=[
RenameFieldStep("cust_id", "customer_id"),
AddFieldStep("customer_segment", compute_fn=lambda row: segment_customer(row)),
AddFieldStep("lifetime_value", default_value=0.0),
RemoveFieldStep("legacy_score"),
CastFieldStep("age", "int32"),
],
)Phase 4: CI/CD Integration
Schema Change Detection in Pull Requests
# .github/workflows/schema-validation.yml
name: Schema Change Validation
on:
pull_request:
paths:
- 'schemas/**'
- 'src/features/**'
- 'pipelines/**'
jobs:
validate-schema-changes:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0 # Full history for diff
- name: Detect schema changes
run: |
git diff origin/main --name-only -- schemas/ > changed_schemas.txt
echo "Changed schemas:"
cat changed_schemas.txt
- name: Validate schema compatibility
run: |
python scripts/validate_schema_changes.py \
--changed-schemas changed_schemas.txt \
--base-branch origin/main
- name: Run data pipeline with new schema
run: |
python scripts/test_pipeline_with_schema.py \
--schema-dir schemas/ \
--test-data data/test/
- name: Check model compatibility
run: |
python scripts/check_model_schema_compatibility.py \
--new-schema schemas/customer_features_v2.json \
--current-model mlflow://models/customer-churn/production
- name: Post PR comment with impact analysis
uses: marocchino/sticky-pull-request-comment@v2
with:
path: schema_change_report.mdConclusion
Schema evolution in ML systems requires a registry-based approach with compatibility enforcement at every pipeline boundary. The schema registry tracks all versions and validates changes against compatibility rules (backward, forward, or full). Runtime validation at ingestion, feature store, and model input boundaries catches schema drift before it causes silent prediction errors. Automated migrations handle field additions, removals, renames, and type changes. CI/CD integration catches incompatible schema changes in pull requests before they reach production. Teams that implement systematic schema management eliminate an entire category of ML pipeline failures—the "it worked yesterday but broke today because someone changed a column name" class of bugs that consume 15-20% of ML engineering time in organizations without schema governance.
