substack note
Recently did panel with Joe Reis Simon Whiteley and the unstoppable Holly Smith at Databricks. We played this fun game called “Room 404: NOT FOUND” where we pick a topic we’d like banished forever and Joe picked the word “context”. We debate and the audience vote
The audience agreed with Joe, and decided to banish the word context to room 404
Well why you ask? Because the definition of “Context” is too broad. What is it? Docs? Lineage? stuff in people’s heads?
You will be pleased to hear that this article’s “context gap” refers to something very specific, so please don’t shoot me for using the word context!!!
Introduction
Recently I’ve been spending some more time in the Bay Area with Data Engineers that are pushing the boundaries of what’s possible, and operating at a scale never seen before.
These teams are doing a few things that we don’t normally see
Outsourcing most, if not all of their work to AI
Speed: focussing more on speed of delivery rather than long-term architecture
Scale: hitting limits of compute availability; not budgetary constraints, but actual compute constraints due to all of AWS’ Capacity / RAM being sold.
The result of this confluence of factors for companies in the bay area is an interesting focus on efficiency for agents. We like the idea of agents doing more work, but it creates a ton of problems:
Agents are good at writing things that start efficient but become inefficient over time; like a SQL Query that works well for small data but deteriorates as volumes increase
Optimising Pipelines at scale is hard: as data surface area changes and different people use and define different metrics, understanding how data is used, stored, and partitioned become imperative for efficient compute utilisation and cost efficiency
Blast Radius becomes imperative: the impact of a change and whether or not it can be done safely is a necessary input for AI-generated code, but this data exists in silo as lineage information is disparate
Recently I spent some time at Open AI’s Headquarters in San Francisco to understand how they were tackling these problems. I should add that I am in no way affiliated with Open AI and don’t claim to speak for their data team, however I can use a few examples that were inspired by some of the excellent talks given by their data engineering teams recently.
o_0
How to build Data Architecture like an AI Lab
Customising Airflow
In 6 weeks, we re-architected and rewrote Apache Airflow’s scheduler, executor, and runner at OpenAI to handle 70,000+ concurrent tasks with 35× lower scheduling latency.
That’s from this post by Ping. OpenAI don’t really useAirflow — they took something with very basic scheduling and customised the shit out of it, including rewriting the scheduler executor and runner in rust.
That’s not a mod, that’s a complete rewrite!
They also rewrote the UI, cause the UI ain’t great.
The point is — Open AI built their own scheduler. As someone that runs a SAAS Company in the space it is a clear epitome of “If you’re not Open AI, don’t act like Open AI”. Not only do Open AI have access to unlimited Codex, but they also have some of the best and expensive data engineers in the world.
A company like Open AI can afford to this. You probably cannot — and it’s probably not justified either, because of the next point I’ll make.
Building for enormous scale
70,000 concurrent tasks is enormous, but so are the numbers involved.
A managed data warehouse like Snowflake or Databricks might make up a tiny <5% portion of an AI Lab’s compute footprint. Open query engines like Spark, Trino and Presto alongside open table formats like iceberg and delta are far more likely.
This means that infrastructure gets built on basic primitives — we’re talking S3 for storage, Kubernetes for compute, and Airflow used how it’s actually meant to be! See here also.
Please don’t run agents or workers on your airflow cluster
This is a huge amount of infrastructure to maintain, but the overhead is justified when you have this much data that needs to be computed efficiently.
Storing Context for AI Agents
Metadata is critical, and there are a few awesome ways you can do this.
DAG State
DAG state answers “what does the pipeline look like right now” — the current shape of every DAG as defined in code, not its execution history.
You parse this directly out of the Airflow repo using Airflow’s own DagBag, the same mechanism the Airflow scheduler uses to load DAG definitions.
Rather than re-implementing DAG parsing or relying on the Airflow metadata DB (which only tells you what has run, not the current definition), we load the repo through DagBag and walk the resulting DAG objects — tasks, dependencies, schedules, operators, params — and serialize the whole structure to JSON.
That JSON becomes the canonical, versioned snapshot of “what this DAG is” at a given commit, decoupled from any particular Airflow deployment’s runtime state.
It’s cheap to regenerate on every repo change and gives every downstream consumer (UI, lineage, alerting) a consistent structural view without needing live access to the Airflow instance.
DAG History
DAG history is the record of what actually happened — every run, every task instance, every state transition, over time.
This lives in Airflow’s own Postgres metadata database, but querying that DB directly at scale (across many customer Airflow instances, or over long time windows) doesn’t hold up.
Instead, we replicate the Postgres DB out of the source Airflow deployment, and sit a Rust-based query server on top of the replica. The Rust layer is what actually serves history queries — it’s built for the query patterns we care about (run history for a DAG, task-level timing, state transitions over a range) and stays fast even as the replicated history grows large.
On top of that query server we expose an MCP so agents and tools can pull DAG history directly — “show me the last 20 runs of this DAG” or “when did this task last fail” — without needing raw DB access or Airflow API round-trips.
Thanks Claude I guess
Asset State
Asset state is metadata about the actual data assets — tables, views, whatever the pipelines produce — rather than the pipelines themselves.
This doesn’t come from Airflow at all; it’s fetched directly from the catalogs of the query engines that sit over the data, e.g. Presto/Trino’s catalog APIs.
That gets us schema, table definitions, partitioning, and other structural metadata straight from the source of truth for the data layer, kept independent of whatever orchestration tool actually produced the asset.
Because it’s pulled from the query engine’s own catalog rather than inferred from pipeline code, it stays accurate even when multiple pipelines or tools write to the same asset.
Asset History
Asset history — how an asset’s state has changed over time, what wrote to it and when — is derived by combining the sources above rather than being its own separate pipeline.
We correlate DAG History (which run touched this asset, when, with what outcome) with periodic snapshots of Asset State (what the asset’s structure looked like at each point) to build a timeline per asset.
In effect, asset history is a join: DAG execution events from the Rust-served history layer, keyed against the asset(s) each task operation touched, giving a per-asset audit trail without maintaining a separate history store.
Lineage
Lineage answers “what does this depend on, and what depends on this.” Rather than storing lineage in a graph database, we take advantage of the fact that the full relation is stored at write time — if A → B → C, the record for C already carries its full upstream chain, not just its immediate parent B.
That means fetching C’s lineage is a single lookup, not a graph traversal: dependencies for both B and A come back with it directly.
This trades a small amount of storage duplication (the same edge information is denormalized into each downstream node) for lineage queries that are effectively O(1) instead of requiring recursive graph walks, which matters a lot at query time for anything user-facing (the lineage UI, impact analysis, alerting on upstream failures).
AI NOTE — read if you’re interesting in parsing lineage from Airflow:
Lineage from an Airflow repo comes from two levels, and it’s worth separating them:
Task-level dependencies (
set_downstream/>>) — these are DAG execution edges, not data lineage. Task A running before Task B doesn’t mean B’s table depends on A’s table.Asset-level lineage — which tables/files a task actually reads (
inlets) and writes (outlets). This is what you actually want for “what does this table depend on,” and Airflow has a built-in (if under-used) mechanism for it: every operator can declareinletsandoutlets, either directly or viaairflow.lineage.AUTOinference from templated fields.
So the extraction path is: load the repo with DagBag (same as DAG State), walk every task, pull its inlets/outlets, and that gives you asset→asset edges through each task. Once you have those edges, you do the transitive-closure step from before: for every asset, compute its full upstream ancestor set, not just its immediate parents, and store that whole set on the node. That’s what makes fetching C return A and B in one lookup — the denormalization happens at build time, not at query time.
Here’s a script that does exactly that:
“”“
Extract asset-level lineage from an Airflow repo.
Approach:
1. Load every DAG in the repo with DagBag (same mechanism used for DAG State).
2. For each task, read its declared `inlets` / `outlets` — Airflow’s built-in
lineage API (https://airflow.apache.org/docs/apache-airflow/stable/
administration-and-deployment/lineage.html). An inlet/outlet is anything
with an identity — most commonly a Table(cluster, database, name) or a
File(url), but any hashable object works.
3. Build a directed graph: for a task with inlets [X] and outlets [Y], add
edge X -> Y (”Y depends on X”).
4. Compute the *transitive closure* per node up front (at build time, not
query time) so that each node’s record already carries its full set of
upstream ancestors. This is the trick that avoids a graph DB: if
A -> B -> C, then fetching C returns {B, A} directly, with no traversal.
5. Serialize to JSON keyed by asset id.
Run:
python extract_lineage.py /path/to/airflow/dags --out lineage.json
“”“
import argparse
import json
import sys
from collections import defaultdict
from airflow.models import DagBag
def asset_id(asset) -> str:
“”“
Normalize an inlet/outlet object to a stable string id.
Airflow’s Table/File lineage objects expose enough fields to build one;
fall back to str() for anything else (e.g. a plain string table name).
“”“
if hasattr(asset, “cluster”) and hasattr(asset, “database”) and hasattr(asset, “name”):
# airflow.lineage.entities.Table-like object
return f”{asset.cluster}.{asset.database}.{asset.name}”
if hasattr(asset, “url”):
# airflow.lineage.entities.File-like object
return asset.url
return str(asset)
def extract_edges(dags_folder: str):
“”“
Walk every DAG/task in the repo and emit (upstream_asset, downstream_asset,
via_dag, via_task) edges from each task’s declared inlets -> outlets.
“”“
dagbag = DagBag(dag_folder=dags_folder, include_examples=False)
if dagbag.import_errors:
print(”Warning: some DAGs failed to import:”, file=sys.stderr)
for path, err in dagbag.import_errors.items():
print(f” {path}: {err}”, file=sys.stderr)
edges = [] # list of (upstream, downstream, dag_id, task_id)
for dag_id, dag in dagbag.dags.items():
for task in dag.tasks:
inlets = getattr(task, “inlets”, None) or []
outlets = getattr(task, “outlets”, None) or []
if not inlets or not outlets:
# A task with no declared inlets/outlets contributes no
# asset-level lineage edge, even if it has upstream/downstream
# *task* dependencies in the DAG.
continue
for src in inlets:
for dst in outlets:
edges.append((asset_id(src), asset_id(dst), dag_id, task.task_id))
return edges
def build_ancestor_closure(edges):
“”“
Build direct-parent adjacency, then compute the full transitive ancestor
set for every node via DFS with memoization. This is the one-time cost
that makes every later lineage fetch O(1).
“”“
parents = defaultdict(set) # node -> immediate upstream parents
provenance = defaultdict(list) # (parent, child) -> [(dag_id, task_id), ...]
nodes = set()
for src, dst, dag_id, task_id in edges:
parents[dst].add(src)
provenance[(src, dst)].append({”dag”: dag_id, “task”: task_id})
nodes.add(src)
nodes.add(dst)
ancestors_cache = {}
def ancestors_of(node, _visiting=None):
if node in ancestors_cache:
return ancestors_cache[node]
_visiting = _visiting or set()
if node in _visiting:
# cyclic lineage shouldn’t happen, but don’t hang if it does
return set()
_visiting.add(node)
result = set()
for parent in parents.get(node, ()):
result.add(parent)
result |= ancestors_of(parent, _visiting)
ancestors_cache[node] = result
return result
closure = {}
for node in nodes:
anc = ancestors_of(node)
closure[node] = {
“immediate_parents”: sorted(parents.get(node, ())),
“all_upstream”: sorted(anc), # <- the denormalized full chain
“upstream_count”: len(anc),
}
return closure, provenance
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(”dags_folder”, help=”Path to the Airflow repo’s dags/ folder”)
parser.add_argument(”--out”, default=”lineage.json”, help=”Output JSON path”)
args = parser.parse_args()
edges = extract_edges(args.dags_folder)
closure, provenance = build_ancestor_closure(edges)
out = {
“assets”: closure,
# kept separately so a lineage lookup stays a flat dict read;
# provenance is only pulled when someone asks “via what task?”
“provenance”: {f”{s}->{d}”: v for (s, d), v in provenance.items()},
}
with open(args.out, “w”) as f:
json.dump(out, f, indent=2, default=str)
print(f”Wrote lineage for {len(closure)} assets, {len(edges)} edges -> {args.out}”)
if __name__ == “__main__”:
main()
# --- Example: querying the output ---
#
# Given DAGs that produce: raw_events -> stg_events -> dim_events
#
# lineage.json[”assets”][”warehouse.analytics.dim_events”] looks like:
#
# {
# “immediate_parents”: [”warehouse.analytics.stg_events”],
# “all_upstream”: [
# “warehouse.analytics.stg_events”,
# “warehouse.raw.raw_events”
# ],
# “upstream_count”: 2
# }
#
# One dict lookup for dim_events gives you stg_events AND raw_events —
# no recursive graph traversal needed at read time.Use-cases for AI Agents
Now something like the Orchestra Runtime really comes in to play.
The Orchestra Runtime is a Harness, Sandbox, and Agentic Orchestrator built for triggering and running complex agent tasks from a single place.
By combining an Orchestrator with an Agentic Harness for Data, you can run workflows on
Non-standard cron intervals
As part of github or CI/CD actions
As part of data pipelines
When certain events happen / sensors
Which is perfect for data team, pipeline maintenance and finops use-cases! Let’s look at a few!
1. Catching SQL and loading patterns that degrade as data grows
A query that’s fine at 10M rows can quietly become the slowest job in the DAG at 500M — and nobody notices until freshness SLAs starts to slip.
Because DAG History already gives you per-task duration over time, and Asset State gives you the underlying table’s row/byte growth from the query engine’s own catalog, you can run a scheduled agent job that joins the two.
An example alert output showing how long something is taking vs. a baseline
For every task, is its duration growing faster than the data it’s processing? A job whose runtime is scaling linearly with data volume is healthy; one scaling worse than linear (a missing partition filter, an accidental cross join, a broken incremental predicate) gets flagged automatically.
Much better than finding out something is taking too long after the fact.
2. Finding the longest path and the bottleneck nodes behind it
Freshness for any given asset is really a function of the slowest chain of dependencies behind it, not any single task. Because DAG State already gives you the dependency graph and DAG History gives you real historical durations per task, an agent can walk the lineage back from any asset and compute the actual critical path.
The hypothetical limiting or longest path of the DAG
This is the sequence of tasks whose cumulative runtime determines when that asset lands — rather than guessing from DAG shape alone.
It surfaces the one or two bottleneck nodes actually worth optimizing, instead of spreading effort across a whole DAG on the assumption that “it’s slow.”
3. Scoring tables by real analytics usage
Information schema and query logs on the engine (Presto/Trino, Snowflake, etc.) already record which tables actually get queried, by whom, and how often.
Layering that usage signal on top of Asset State turns “we have 4,000 tables” into a ranked list of which tables genuinely drive analytics versus which are stale, orphaned, or only ever touched by the pipeline that built them.
Example metadata helping AI to decide what customers table to use
Run this periodically via your Runtime layer and you get a live “most valuable tables” view — useful for prioritizing documentation, SLAs, and access cleanup without anyone manually auditing the warehouse.
4. Inferring how a table should be partitioned from its own column make-up
Where the data in a table actually concentrates — which column(s) it clusters around, how skewed the distribution is — is the single best signal for how that table should be partitioned or clustered, and it’s usually decided once at table-creation time and never revisited.
For example, you might have a bigquery table orders partitioned by date and by customer . You might see that 80% of queries go like
select
*
from orders
where date = ‘2026-01-01’
and product_Id = ‘ipads’You could then change the partitioning to be based on product_id instead of customer — this saves time, money etc.
5. Scoring the blast radius of a PR before it merges
This is where lineage and usage data combine directly. A PR that changes a model or a DAG touches some asset; that asset’s downstream dependents are already known instantly by walking the same denormalized closure used for upstream lineage.
It’s sort of like a combination of the other ones, because the output is something like:
{
“asset”: “analytics.dim_customers”,
“change”: “PR #4821: drop column `legacy_segment_id`”,
“downstream_dependencies”: [
{
“name”: “analytics.fct_customer_orders”,
“type”: “table”,
“last_7_days_queries”: 8460,
“size_gb”: 14.2,
“size_rows”: 31200000,
“tags”: [”core”, “finance”, “daily-refresh”]
},
{
“name”: “analytics.rpt_customer_ltv”,
“type”: “table”,
“last_7_days_queries”: 1120,
“size_gb”: 3.6,
“size_rows”: 4800000,
“tags”: [”finance”, “pii”]
},
{
“name”: “Exec Revenue Overview”,
“type”: “dashboard”,
“last_7_days_queries”: 640,
“size_gb”: null,
“size_rows”: null,
“tags”: [”exec-dashboard”, “daily-refresh”]
},
{
“name”: “analytics.dim_customers_regional”,
“type”: “table”,
“last_7_days_queries”: 95,
“size_gb”: 1.8,
“size_rows”: 4750000,
“tags”: [”regional”, “unowned”]
},
{
“name”: “Churn Risk Monitor”,
“type”: “dashboard”,
“last_7_days_queries”: 12,
“size_gb”: null,
“size_rows”: null,
“tags”: [”customer-success”, “stale”]
}
],
“blast_radius_summary”: {
“total_dependents”: 5,
“tables”: 3,
“dashboards”: 2,
“high_usage_dependents”: 2,
“flagged_tags”: [”exec-dashboard”, “pii”, “finance”]
}
}This is something that could be fed back to a PR, to tag it as HIGH_RISK or perhaps LOW_RISK , which in turn could inform how and who reviews PRs.
Conclusion — plugging the context gap with metadata
It’s clear that while large companies like Open AI probably have the budget and incentives to build customisable architecture, having these agents is something anyone can benefit from.
After all, who wouldn’t want an AI engineer looking out for them, to help tweak nobs, identify anomalies, and fundamentally stop stuff breaking without any work?
The barrier is the Context Gap. The Context Gap refers to this treasure trove of metadata that is disparate across systems.
This is only going to get more complex. As companies get more complex and shift to multi-query engines, context will be more sparsely distributed.
The Context Gap
In enterprises and with AI, human-driven and AI-generated context is both expanding and becoming more disparate. This is driving Data Platform Teams to explore Agentic Control Planes like Orchestra — where a managed, agentic control plane provides this orchestration, runtime, and context out the box.
Orchestra intelligently gathers and structures metadata from different sources into a centralised Context Repository. Orchestra provides the Agent Runtime (=Harness+Sandbox), a Data-specific Harness and Sandbox capable of finding the right ways to optimise pipelines, and provides an orchestrator for scheduling and running these tasks.
Hopefully this article gives an idea for what the options are for plugging the context gap, all the way from managed solutions like Orchestra to going the full-blown build-it-yourself route, just like the AI Labs.
For more content like this, please like, subscribe, and consider following on other platforms!