Dataset Lineage and Management
The 12% That Wasn't There
The computer vision team is thrilled. Their new image classifier is achieving 94.7% accuracy on the test set - up from 82.3% six months ago. The 12.4% improvement justifies the entire quarter of experimentation. They present the results at the all-hands meeting. The model is approved for production deployment.
Three weeks later, the model is live. Actual production accuracy: 81.9%. Barely better than the model it replaced.
What happened?
A post-mortem investigation reveals the root cause: the augmentation pipeline that generated the training data was applied before the train/test split. One of the augmentations - a random crop - created multiple variants of each image. An image that appeared in the "test set" also had its augmented siblings in the training set. The model had effectively seen the test images (or very similar versions of them) during training.
The 12.4% accuracy improvement was 98% leakage, 2% real improvement.
This is test set leakage. It is one of the most damaging and hardest-to-detect bugs in ML. Understanding and preventing it - and building the lineage tracking to detect it after the fact - is what this lesson covers.
:::tip 🎮 Interactive Playground Visualize this concept: Try the Dataset Lineage & Provenance demo on the EngineersOfAI Playground - no code required. :::
What Is Dataset Lineage
Dataset lineage is the complete record of a dataset's provenance: where it came from, what transformations were applied, and how it relates to other datasets. A lineage graph answers:
- Which raw source files contributed to this training dataset?
- What preprocessing and feature engineering steps were applied, in what order, with what parameters?
- Which training dataset version trained which model?
- If this dataset was created by splitting another dataset, what is the relationship between the splits?
- When was each transformation applied, and who triggered it?
Lineage tracking would have flagged this: a lineage system that records "augmentation applied before split" would allow the team to query "do any test set examples share a source image with any training set example?" and catch the leakage before the model went to production.
Types of Data Leakage
Leakage is any case where information from outside the training set contaminates the training process, producing overly optimistic evaluation metrics.
1. Temporal Leakage
Using future information to predict past events. Example: predicting whether a user will churn next month using features computed after the churn event.
# WRONG: feature computed after the label event
df_training = (
df_events
.join(df_user_features, on="user_id") # features might be from after the event
.withColumn("label", F.col("churned_within_30d"))
)
# CORRECT: point-in-time correct join
df_training = (
df_events
.join(
df_user_features,
on=(
(df_events.user_id == df_user_features.user_id) &
# Only use features that existed BEFORE the label event
(df_user_features.feature_date < df_events.event_date)
)
)
.withColumn("label", F.col("churned_within_30d"))
)
2. Augmentation Leakage
Applying data augmentation before splitting creates related examples across splits. Any augmentation that produces variants of a single source example must be applied after splitting and only to the training split.
import numpy as np
from sklearn.model_selection import train_test_split
def create_correct_splits(
image_paths: list,
labels: list,
test_size: float = 0.2,
val_size: float = 0.15,
seed: int = 42,
) -> dict:
"""
Create train/val/test splits BEFORE any augmentation.
Returns splits by SOURCE file - augmentation happens later.
"""
# Split on source files, not on individual examples
train_val_paths, test_paths, train_val_labels, test_labels = train_test_split(
image_paths, labels, test_size=test_size, random_state=seed,
stratify=labels, # maintain class distribution
)
val_size_adjusted = val_size / (1 - test_size)
train_paths, val_paths, train_labels, val_labels = train_test_split(
train_val_paths, train_val_labels,
test_size=val_size_adjusted, random_state=seed,
stratify=train_val_labels,
)
return {
"train": {"paths": train_paths, "labels": train_labels},
"val": {"paths": val_paths, "labels": val_labels},
"test": {"paths": test_paths, "labels": test_labels},
}
# CORRECT usage:
splits = create_correct_splits(all_image_paths, all_labels)
# Now apply augmentation ONLY to the training split
def augment_training_data(train_paths, train_labels, n_augments=4):
"""Apply augmentation to training split only, after splitting."""
augmented_paths = []
augmented_labels = []
for path, label in zip(train_paths, train_labels):
# Original
augmented_paths.append(path)
augmented_labels.append(label)
# Augmented variants (only generated from training images)
for i in range(n_augments):
aug_path = apply_augmentation(path, seed=i)
augmented_paths.append(aug_path)
augmented_labels.append(label)
return augmented_paths, augmented_labels
train_paths_aug, train_labels_aug = augment_training_data(
splits["train"]["paths"], splits["train"]["labels"]
)
3. Preprocessing Leakage
Fitting preprocessing transformations (scalers, encoders) on the full dataset before splitting.
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
import numpy as np
X = np.random.randn(1000, 10)
y = np.random.randint(0, 2, 1000)
# WRONG: fit scaler on all data, then split
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X) # uses test data statistics!
X_train, X_test = X_scaled[:800], X_scaled[800:]
# CORRECT: split first, fit scaler only on training data
X_train_raw, X_test_raw = X[:800], X[800:]
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train_raw) # fit only on train
X_test = scaler.transform(X_test_raw) # transform with train statistics
4. ID Leakage (Group Leakage)
When examples are not independent - e.g., multiple samples from the same patient, multiple reviews from the same user - splitting randomly will put related examples in both train and test sets.
from sklearn.model_selection import GroupShuffleSplit
# WRONG: random split may put patient A's scans in both train and test
X_train, X_test = train_test_split(X, test_size=0.2)
# CORRECT: group-aware split - keep all of patient A's data together
patient_ids = df["patient_id"].values
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
train_idx, test_idx = next(gss.split(X, y, groups=patient_ids))
X_train, X_test = X[train_idx], X[test_idx]
y_train, y_test = y[train_idx], y[test_idx]
# Verify no patient appears in both splits
train_patients = set(patient_ids[train_idx])
test_patients = set(patient_ids[test_idx])
overlap = train_patients & test_patients
assert len(overlap) == 0, f"Leakage: {len(overlap)} patients in both splits"
Stratified Splitting
Stratification ensures that the class distribution in each split matches the overall distribution. Without it, small classes may be over- or under-represented in splits, causing unstable evaluation metrics.
import pandas as pd
from sklearn.model_selection import StratifiedShuffleSplit
from collections import Counter
def stratified_split(
df: pd.DataFrame,
label_col: str,
val_size: float = 0.15,
test_size: float = 0.20,
seed: int = 42,
) -> tuple:
"""
Create stratified train/val/test splits.
Returns DataFrames with split column added.
"""
# First: split off test set
sss1 = StratifiedShuffleSplit(n_splits=1, test_size=test_size, random_state=seed)
train_val_idx, test_idx = next(sss1.split(df, df[label_col]))
df_train_val = df.iloc[train_val_idx].copy()
df_test = df.iloc[test_idx].copy()
# Second: split train_val into train and val
val_size_adjusted = val_size / (1 - test_size)
sss2 = StratifiedShuffleSplit(
n_splits=1, test_size=val_size_adjusted, random_state=seed + 1
)
train_idx, val_idx = next(sss2.split(df_train_val, df_train_val[label_col]))
df_train = df_train_val.iloc[train_idx].copy()
df_val = df_train_val.iloc[val_idx].copy()
# Add split labels
df_train["split"] = "train"
df_val["split"] = "val"
df_test["split"] = "test"
# Verify class distribution is preserved
print("Class distribution:")
print(f" Overall: {Counter(df[label_col])}")
print(f" Train: {Counter(df_train[label_col])}")
print(f" Val: {Counter(df_val[label_col])}")
print(f" Test: {Counter(df_test[label_col])}")
return df_train, df_val, df_test
# Leakage detection: verify no example appears in multiple splits
def verify_no_leakage(
df_train: pd.DataFrame,
df_val: pd.DataFrame,
df_test: pd.DataFrame,
id_col: str = "example_id",
):
"""Assert no data leakage between splits."""
train_ids = set(df_train[id_col])
val_ids = set(df_val[id_col])
test_ids = set(df_test[id_col])
assert len(train_ids & val_ids) == 0, "Leakage: train-val overlap"
assert len(train_ids & test_ids) == 0, "Leakage: train-test overlap"
assert len(val_ids & test_ids) == 0, "Leakage: val-test overlap"
total = len(train_ids) + len(val_ids) + len(test_ids)
all_ids = set(df_train[id_col]) | set(df_val[id_col]) | set(df_test[id_col])
assert total == len(all_ids), "Some examples appear in multiple splits"
print(f"No leakage detected. Total examples: {total:,}")
print(f" Train: {len(train_ids):,} ({100*len(train_ids)/total:.1f}%)")
print(f" Val: {len(val_ids):,} ({100*len(val_ids)/total:.1f}%)")
print(f" Test: {len(test_ids):,} ({100*len(test_ids)/total:.1f}%)")
Dataset Registries
A dataset registry is a catalog of all datasets your team uses, with metadata for discoverability and governance. It answers: "What datasets exist? Who owns them? What models use them? What is the latest valid version?"
from dataclasses import dataclass, field
from typing import Optional
from datetime import datetime
import json
@dataclass
class DatasetCard:
"""Standard dataset documentation card."""
name: str
version: str
description: str
owner_team: str
owner_contact: str
created_at: str
# Location
storage_path: str
format: str # parquet, delta, csv, etc.
# Content
record_count: int
size_bytes: int
features: list[str]
label_column: Optional[str]
label_classes: Optional[list]
# Quality
missing_value_pct: float
class_imbalance_ratio: Optional[float]
date_range: Optional[tuple]
# Lineage
source_datasets: list[str] = field(default_factory=list)
preprocessing_script: Optional[str] = None
preprocessing_git_sha: Optional[str] = None
# Governance
contains_pii: bool = False
gdpr_jurisdiction: Optional[str] = None
data_license: Optional[str] = None
approved_use_cases: list[str] = field(default_factory=list)
retention_days: int = 365
# Status
status: str = "active" # active | deprecated | archived
def to_dict(self) -> dict:
return {k: v for k, v in self.__dict__.items()}
def validate(self):
"""Validate the dataset card has required fields."""
assert self.owner_contact, "owner_contact is required"
assert self.record_count > 0, "record_count must be positive"
assert self.format in ["parquet", "delta", "csv", "tfrecord", "jsonl"], \
f"Unknown format: {self.format}"
if self.contains_pii:
assert self.gdpr_jurisdiction, "gdpr_jurisdiction required when contains_pii=True"
return True
# Example registry entry
user_features_card = DatasetCard(
name="user_behavioral_features",
version="v2024q3_002",
description=(
"Behavioral features for all active users computed from clickstream data. "
"Includes 30-day session counts, 90-day purchase history, and segmentation labels."
),
owner_team="data-platform",
created_at="2024-10-01T09:00:00Z",
storage_path="s3://ml-data/delta/user_behavioral_features/",
format="delta",
record_count=15_234_891,
size_bytes=28_912_345_678,
features=["user_id", "session_count_30d", "purchase_amount_90d", "user_segment",
"last_active_date", "device_type", "geo_region"],
label_column=None,
label_classes=None,
missing_value_pct=2.3,
class_imbalance_ratio=None,
date_range=("2024-01-01", "2024-09-30"),
source_datasets=["raw_clickstream_2024q3", "user_profiles_v7"],
preprocessing_script="pipelines/feature_engineering/user_features.py",
preprocessing_git_sha="f7b3a1c9",
contains_pii=True,
gdpr_jurisdiction="EU",
data_license="internal",
approved_use_cases=["recommendation_model", "churn_prediction", "ad_targeting"],
retention_days=180,
status="active",
)
user_features_card.validate()
# Save to a registry (could be a database, git repo, or catalog)
with open("registry/user_behavioral_features_v2024q3_002.json", "w") as f:
json.dump(user_features_card.to_dict(), f, indent=2)
OpenLineage: Standardized Lineage Collection
OpenLineage is an open standard for collecting lineage events from data pipelines. It defines a common API that pipeline tools (Apache Airflow, dbt, Spark, Flink) can emit lineage events to, and that downstream systems (Marquez, DataHub, OpenMetadata) can consume.
from openlineage.client import OpenLineageClient
from openlineage.client.run import (
RunEvent, RunState, Run, Job, Dataset,
InputDataset, OutputDataset,
)
from openlineage.client.uuid import generate_new_uuid
from datetime import datetime, timezone
client = OpenLineageClient(url="http://marquez:5000")
# Emit a START event when the training job begins
run_id = generate_new_uuid()
client.emit(
RunEvent(
eventType=RunState.START,
eventTime=datetime.now(timezone.utc).isoformat(),
run=Run(runId=run_id),
job=Job(namespace="ml-training", name="ctr_model_training_v7"),
inputs=[
InputDataset(
namespace="s3://ml-data",
name="delta/user_behavioral_features",
facets={
"datasetVersion": {"version": "v2024q3_002"},
"columnLineage": {
"fields": {
"session_count_30d": {"inputFields": [
{"namespace": "s3://ml-data",
"name": "raw/clickstream_2024q3",
"field": "session_id"}
]}
}
}
},
),
],
outputs=[],
)
)
# ... training ...
# Emit COMPLETE event with the model as output
client.emit(
RunEvent(
eventType=RunState.COMPLETE,
eventTime=datetime.now(timezone.utc).isoformat(),
run=Run(runId=run_id),
job=Job(namespace="ml-training", name="ctr_model_training_v7"),
inputs=[
InputDataset(namespace="s3://ml-data",
name="delta/user_behavioral_features"),
],
outputs=[
OutputDataset(
namespace="mlflow://tracking",
name="models/ctr_ranker",
facets={
"outputStatistics": {
"rowCount": num_training_samples,
"size": model_size_bytes,
}
},
),
],
facets={
"jobMetrics": {
"val_auc": 0.8912,
"training_time_seconds": 14400,
}
},
)
)
Lineage-Based Leakage Detection
With lineage tracking, you can write automated checks that detect leakage patterns:
import hashlib
import pandas as pd
import numpy as np
def detect_augmentation_leakage(
train_df: pd.DataFrame,
test_df: pd.DataFrame,
source_id_col: str = "source_image_id",
) -> dict:
"""
Detect if test set contains augmented variants of training images.
Requires a source_id_col that maps augmented images to their source.
"""
train_source_ids = set(train_df[source_id_col].unique())
test_source_ids = set(test_df[source_id_col].unique())
leaked_sources = train_source_ids & test_source_ids
return {
"has_leakage": len(leaked_sources) > 0,
"leaked_source_count": len(leaked_sources),
"leaked_test_examples": int(
test_df[test_df[source_id_col].isin(leaked_sources)].shape[0]
),
"leakage_pct_test": 100 * len(
test_df[test_df[source_id_col].isin(leaked_sources)]
) / len(test_df),
"example_leaked_ids": list(leaked_sources)[:5],
}
def compute_feature_target_correlation(
df: pd.DataFrame,
feature_cols: list,
target_col: str,
threshold: float = 0.95,
) -> list:
"""
Flag features with suspiciously high correlation to target.
High correlation may indicate target leakage (feature computed using label).
"""
suspicious = []
for col in feature_cols:
if df[col].dtype in [float, int]:
corr = abs(df[col].corr(df[target_col]))
if corr > threshold:
suspicious.append({
"feature": col,
"correlation": corr,
"flag": "potential_target_leakage",
})
return suspicious
Sharding and Partitioning for Large Datasets
For very large training datasets, you need a partitioning strategy that enables efficient loading during training.
from pyspark.sql import functions as F
# Partition a large dataset for efficient ML loading
def create_ml_shards(
spark,
df,
output_path: str,
n_shards: int = 256,
target_shard_size_mb: int = 128,
) -> dict:
"""
Create evenly-sized shards with random shuffling.
Each shard is balanced across classes.
"""
# Add a random partition key for balanced sharding
df_sharded = df.withColumn(
"shard_id",
(F.rand(seed=42) * n_shards).cast("int")
)
# Write shards as parquet files
(
df_sharded
.repartition(n_shards, "shard_id")
.write
.format("parquet")
.mode("overwrite")
.option("maxRecordsPerFile", 100_000)
.save(output_path)
)
# Record shard metadata
shard_stats = (
df_sharded
.groupBy("shard_id")
.agg(
F.count("*").alias("record_count"),
F.countDistinct("label").alias("n_classes"),
)
.toPandas()
)
return {
"n_shards": n_shards,
"total_records": int(df.count()),
"output_path": output_path,
"shard_stats": shard_stats.to_dict("records"),
}
Common Mistakes
:::danger Applying Any Preprocessing Before the Split Normalization, encoding, augmentation, resampling - any preprocessing that looks at distribution statistics or creates derived examples must happen after the train/val/test split. The only exception: deterministic, per-example transformations that do not use any global statistics (e.g., resizing an image to a fixed size, lowercasing text). :::
:::danger Using Test Set Leakage Detection Only at the End Leakage detection should run automatically as part of your data pipeline, not as a one-time post-mortem. Add a leakage check step to your DVC pipeline or Airflow DAG that verifies split integrity every time the dataset is regenerated. :::
:::warning Splitting by Row Index Instead of by Entity In datasets where one entity (user, patient, store) appears multiple times, random row-based splitting puts the same entity in both train and test. Use group-aware splitting (GroupShuffleSplit or similar) whenever your dataset has a natural grouping structure. :::
:::warning Not Recording Split Indices The split indices (which records went to train, val, test) are an artifact that must be versioned alongside the data. Without them, two engineers using "the same dataset" may use different splits, making their experiment results incomparable. :::
Interview Q&A
Q: What is temporal leakage and how do you prevent it?
A: Temporal leakage occurs when features used to train a model contain information that would not have been available at prediction time in production. The classic example: predicting user churn using a feature like "total purchases ever" - this feature grows over time and includes purchases made after the prediction date. Prevention: use point-in-time correct joins where features are computed as of a specific timestamp (always before the label event). For time series data, always split chronologically (all training data before a cutoff date, all test data after) rather than randomly. Test for temporal leakage by checking if any feature has a suspiciously high correlation with the label - this often indicates the feature was computed using information that includes the label.
Q: Explain the difference between data lineage and data provenance.
A: Data provenance is the historical record of a dataset's origin - where the raw data came from (source systems, collection methods, dates). Data lineage extends this to include all transformations applied to the data from its source to its current form. Think of provenance as "where did this come from?" and lineage as "how did it get here?" Both are required for full auditability. For ML, lineage is more operationally critical - it lets you trace a model's behavior back to specific data transformations and identify where a bug or distribution shift was introduced.
Q: How would you build a system to detect data leakage in a large production ML pipeline?
A: Four layers: (1) ID-level check - verify that no entity ID (user, patient, item) appears in both training and test splits. Implemented as a pipeline step that computes set intersection and fails if non-empty. (2) Source-level check - for derived datasets (augmentation, oversampling), verify that source IDs of test examples do not overlap with training examples. (3) Feature-target correlation check - flag features with suspiciously high correlation to the label (above 0.95) as potential target leakage candidates for human review. (4) Temporal consistency check - for time-series data, verify that all feature timestamps in the test set are after the last training timestamp.
Q: What is a dataset card and why does it matter for ML teams?
A: A dataset card is a structured documentation artifact that describes a dataset - its content, provenance, quality characteristics, governance properties (PII, license, approved uses), and lineage. It is the "README" for a dataset. Dataset cards matter because: (1) they enable discoverability - engineers can find existing datasets without recreating them; (2) they document approved use cases - preventing a dataset collected for one purpose from being misused for another; (3) they capture data quality issues (missing value rates, class imbalance) that affect model behavior; (4) they document PII and retention requirements for compliance. The Hugging Face model card and dataset card format has become a de facto standard in the ML community.
Q: How do you handle dataset versioning for a feature store that is updated daily?
A: Use Delta Lake or Iceberg as the underlying storage format, which provides time travel natively. Each daily update creates a new Delta version. Training jobs pin to a specific version: the version number is logged in the experiment tracker alongside the run, creating a permanent link between the model and the exact data it was trained on. For the feature store, never use overwrite mode - use MERGE (upsert) to update existing records and add new ones. This preserves full history. Set the VACUUM retention to at least 90 days so that any model trained in the last 90 days can have its training data reconstructed. For GDPR compliance, implement row-level deletes using Delta Lake's DELETE statement, and retrain any model that used the deleted records.