Track Data Lineage

Picture this: your AI pipeline just produced a suspicious prediction, your model's accuracy dropped overnight, or a data scientist asks, "Where did this number actually come from?" — and you have no idea. Without tracking data lineage, debugging a pipeline is like fixing a car with the hood welded shut. Every transformation, join, or aggregation that silently alters your data becomes a black box. In applied AI engineering, tracking data lineage is not a nice-to-have; it's the difference between trust and guesswork. By the end of this lesson, you'll know exactly how to trace every field from raw source to final model output — and you'll have a working implementation to prove it.

The problem this lesson solves

Modern AI pipelines are not linear. They ingest from multiple sources, apply cleansing, feature engineering, and validation, and then feed a model that produces predictions. Somewhere in that chain, a date column gets parsed incorrectly, a join duplicates rows, or a feature scaling step shifts distributions. When something goes wrong, you need to answer two questions fast: What changed? and Where did it happen?

Without lineage, you're stuck grepping logs, reading commit history, and manually retracing steps — a process that takes hours and still misses edge cases. In regulated industries (finance, healthcare, insurance), you also have compliance requirements: auditors demand proof of where each data point came from and how it was transformed. Manually documenting this is error-prone and scales terribly.

Lineage tracking turns your pipeline from a chaotic flow into an auditable, debuggable map. It lets you:

• Reproduce results — know the exact input versions and transformations that produced an output.
• Debug faster — pinpoint which step introduced a bug or anomaly.
• Ensure compliance — provide a clear trail for auditors and stakeholders.
• Improve data quality — identify which sources are unreliable.

This lesson gives you a pragmatic, hands-on approach to adding lineage to your Python-based AI pipelines — without a heavy enterprise tool if you don't need one yet.

Core concept / mental model

Think of data lineage as the genealogy of your data. Just as a family tree shows ancestors and relationships, lineage shows the origin of every column, row, or dataset, plus every operation that transformed it between source and consumption.

In a pipeline, you can capture lineage at two levels:

• Table/Dataset level — tracks which datasets feed into which other datasets (e.g., raw_events → cleaned_events → feature_store).
• Field/Column level — tracks individual columns (e.g., user.age → age_zscore → model_input_age).

A mental model: lineage as metadata. Every time your code reads or writes data, you append a small record: input name, output name, timestamp, transformation type, code version, and any parameters. Over time, these records form a directed acyclic graph (DAG). You can then query that graph to answer "what is upstream of this dataset?" or "what downstream datasets are affected by this source change?"

Keep it simple at first — you don't need a full enterprise solution. A lightweight lineage tracker can be as simple as a decorator or a context manager that wraps your data transformations. The key is consistency: every step must record its lineage for the graph to be useful.

How it works step by step

Here's the standard sequence to implement lineage tracking in a pipeline:

1. Define a lineage schema — decide what information you need per step. At minimum: source, destination, timestamp, operation, code_version, and parameters.
2. Instrument your pipeline steps — wrap each read/transform/write with lineage capture. In Python, you can use decorators, context managers, or explicit calls.
3. Store lineage records — persist to a simple store (JSON lines, SQLite, or a dedicated lineage catalog).
4. Expose a query interface — write functions to answer "upstream" and "downstream" queries.
5. Integrate with existing tools — if you use Airflow or Dagster, they have built-in hooks; but a custom tracker works for any pipeline.

Cause and effect: without step 1 (schema), you'll capture inconsistent data; without step 2 (instrumentation), you'll have gaps; without step 3 (storage), you'll lose history. Each step builds on the previous one.

Hands-on walkthrough

Let's build a minimal lineage tracker in Python. We'll use a decorator to automatically record lineage for functions that transform data.

Step 1: Define a lineage record

First, create a dataclass to hold lineage information.

from dataclasses import dataclass, asdict
import uuid
from datetime import datetime, timezone

@dataclass
class LineageRecord:
    node_id: str
    source: str
    destination: str
    operation: str
    timestamp: str
    code_version: str
    parameters: dict
    parent_ids: list

Step 2: Build a tracker

Now, implement a simple tracker that stores records in memory and can write to JSON Lines.

import json
from pathlib import Path

class LineageTracker:
    def __init__(self, store_path: str = "lineage.jsonl"):
        self.store_path = store_path
        self.records = []

    def record(self, record: LineageRecord):
        self.records.append(record)
        # Append to JSONL file
        with open(self.store_path, "a") as f:
            f.write(json.dumps(asdict(record)) + "\n")

    def upstream(self, node_id: str) -> list:
        """Return parent node ids for a given node."""
        return [r.parent_ids for r in self.records if r.node_id == node_id]

    def downstream(self, node_id: str) -> list:
        """Return child node ids for a given node."""
        children = []
        for r in self.records:
            if node_id in r.parent_ids:
                children.append(r.node_id)
        return children

Step 3: Create a lineage-aware decorator

Wrap your transformation functions with a decorator that records lineage automatically.

def lineage(source: str, destination: str, operation: str, tracker: LineageTracker, parents: list = None):
    def decorator(func):
        def wrapper(*args, **kwargs):
            result = func(*args, **kwargs)
            record = LineageRecord(
                node_id=str(uuid.uuid4()),
                source=source,
                destination=destination,
                operation=operation,
                timestamp=datetime.now(timezone.utc).isoformat(),
                code_version="1.0.0",  # in practice, use your VCS version
                parameters=kwargs,
                parent_ids=parents or []
            )
            tracker.record(record)
            return result, record.node_id
        return wrapper
    return decorator

Step 4: Apply to a sample pipeline

Let's test it with a simple transformation chain.

tracker = LineageTracker()

@lineage(source="raw.csv", destination="cleaned.csv", operation="clean", tracker=tracker)
def clean_data(df, **kwargs):
    # imagine real cleaning logic
    return df.dropna()

@lineage(source="cleaned.csv", destination="features.parquet", operation="feature_eng", tracker=tracker, parents=["clean"])
def engineer_features(df, **kwargs):
    # imagine feature creation
    return df

cleaned, clean_id = clean_data(pd.DataFrame({"a": [1, None, 3]}))
features, feat_id = engineer_features(cleaned, min_vals=True)

# Query lineage
print(tracker.upstream(feat_id))  # Should print parent ids
print(tracker.downstream(clean_id))  # Should print the feature node id

Expected output:

[["<uuid-for-clean>"] or [] if parents not passed properly]
["<uuid-for-features>"]

Pro tip: Keep code_version meaningful — use git rev-parse HEAD or a CI-provided version to tie lineage to exact code.

Compare options / when to choose what

You don't always need a custom tracker. Here's a comparison of common approaches:

When to choose what:

• If you're prototyping or have a handful of pipelines, custom tracker is a fast win.
• If you need formal data contracts and alerting, Great Expectations is a good complement.
• If your pipelines already run on Airflow or Dagster, use their built-in lineage features before adding external tools.
• For enterprise-scale governance, invest in a data catalog like DataHub or Amundsen — they handle multi-team collaboration and compliance at scale.

Troubleshooting & edge cases

• Missing lineage records: Make sure every step that reads or writes data is wrapped with your @lineage decorator. If your pipeline uses external tools (SQL, Spark), you'll need to add explicit logging.
• Circular references: If a pipeline step reads from a dataset it also writes to, you can create a cycle in your lineage graph. Always read from an intermediate version or use a versioned store.
• Branching/merging: When multiple parents feed one child, store parent_ids as a list, not a scalar. Your upstream() function should return all parents, not just one.
• Parameter serialization: If you store parameters containing DataFrames or non-JSON-serializable objects, you'll crash. Convert them to safe types (e.g., strings, numbers, dicts) before saving.
• Timestamp consistency: Use UTC ISO timestamps to avoid timezone confusion across environments.
• Code version accuracy: If code_version is hardcoded, it becomes stale. Use importlib.metadata.version or a build-time constant.

What you learned & what's next

You've now mastered the core of tracking data lineage in pipelines — you know why it's essential, how to model it as a graph, and how to implement a lightweight tracker in Python. You can instrument your own transformation functions, store lineage records, and query upstream/downstream dependencies. This puts you far ahead of most practitioners who rely on guesswork.

You've achieved the lesson objectives: you can explain the concept and complete a practical exercise. Remember these key points:

• Lineage is about metadata that ties data through transformations.
• A simple decorator-based tracker is enough for many pipelines.
• Store lineage as a graph; query it for debugging and compliance.

What's next in your Applied AI engineering path? You'll likely learn about pipeline observability or model monitoring — but armed with lineage, you can debug any data issue with confidence. Keep this tracker handy; you'll reuse it in upcoming lessons.

Final pro tip: Start small — add lineage to one critical pipeline first. Once you see how it accelerates debugging, you'll want it everywhere.