Column Lineage section silently vanishes from two-stage prompts when the pipeline is deeper than lineage_expansion_depth

[mingjerli]

Version: 0.0.8

build_lineage_context() returns "" — dropping the entire ## Column Lineage section from the two-stage text2sql prompt — for any pipeline whose lineage is deeper than ContextConfig.lineage_expansion_depth (default 2). It is all-or-nothing rather than degraded: one extra level of depth flips it from zero lines to complete.

Reproduction

A four-level pipeline (source → raw → stg → mart), which is an ordinary warehouse shape:

from clgraph import Pipeline
from clgraph.tools.context import ContextBuilder, ContextConfig

SQL = [
    ("raw",  "CREATE TABLE raw.orders AS SELECT order_id, amount FROM source.orders"),
    ("stg",  "CREATE TABLE stg.orders AS SELECT order_id, amount FROM raw.orders"),
    ("mart", "CREATE TABLE mart.daily AS SELECT SUM(amount) AS revenue FROM stg.orders"),
]
p = Pipeline(SQL, dialect="bigquery")

for depth in (2, 3):
    b = ContextBuilder(p, ContextConfig(lineage_expansion_depth=depth))
    tables = b.resolve_context_tables(b.expand_with_lineage(["mart.daily"]))
    print(depth, tables)
    print("  ", repr(b.build_lineage_context(tables)))

2 ['mart.daily', 'stg.orders', 'raw.orders']
   ''
3 ['mart.daily', 'stg.orders', 'raw.orders', 'source.orders']
   '## Column Lineage\n\n- mart.daily.revenue <- source.orders.amount\n- stg.orders.order_id <- source.orders.order_id\n- ...'

At the default depth the section is empty. At depth 3 it is five complete lines.

Mechanism

Two functions measure lineage differently, and build_lineage_context requires them to agree:

• expand_with_lineage() walks ancestor tables, bounded by lineage_expansion_depth.
• trace_column_backward() returns terminal source columns only — not the intermediate hops.
• build_lineage_context() then keeps a line only if a traced source is inside the expanded set:

sources = self.pipeline.trace_column_backward(table_name, col.column_name)
relevant_sources = [s for s in sources if s.table_name in tables]

if relevant_sources and relevant_sources[0].table_name != table_name:
    ...

So the two only intersect when the expansion happens to reach all the way to the leaves. In the repro, trace_column_backward("mart.daily", "revenue") returns exactly one entry, source.orders.amount, which sits at depth 3 — one level outside the depth-2 set. Every line is therefore filtered out.

Worth noting: stg.orders.amount and raw.orders.amount are in the expanded set, and would be the most useful lines to show a model reasoning over a bounded table list — but they are never returned by the trace at all, so the filter cannot match them even in principle.

Why it matters

For a pipeline deeper than the default, _generate_two_stage builds its prompt with no lineage section at all. Nothing surfaces this: the call succeeds, the section is simply absent, and the model loses the column-derivation context that 0.0.8 added. The deeper the warehouse, the more certain the loss — which is the opposite of the intent.

This is reachable at defaults on realistic shapes. A source → raw → staging → mart pipeline is four levels, so the default depth of 2 never reaches the leaves.

Possible directions

Not a patch suggestion — the right fix depends on what the section is meant to say:

1. Cite the nearest in-set ancestor instead of requiring the terminal source to be in the set. This seems closest to the intent: the prompt lists a bounded set of tables, so lineage between those tables is the useful statement, and it makes the section degrade gracefully rather than vanish.
2. Let it cite out-of-set sources, dropping the relevant_sources filter and labelling them as upstream of the context.
3. Raise the default depth so expansion reaches the leaves. Simplest, but only moves the cliff — it reappears one layer deeper, and widens the table set for every prompt.

Happy to send a PR for whichever direction you prefer.

How this surfaced

Found while building clgraph Studio's Ask panel, which displays the context clgraph assembles for a prompt. It reproduces _generate_two_stage's config exactly, so the empty section showed up as a visible gap next to a populated schema, relationships and join-hints section.