grad on DataFusion: optimize_unions fails on a nested union of tables that aren't wrt #236

😈🧪 Adversarial tester

Since #229, grad on DataFusion fails on some unions that nest a union of tables that aren't wrt:

DataFusion error: Optimizer rule 'optimize_unions' failed
caused by
Schema error: No field named x. Valid fields are "a.x - Float64(0.5)".

Repro (main at a223f33, ddxdb built from it). b is the wrt table, and a is data:

import pyarrow as pa, ddxdb.ad as ad
from datafusion import SessionContext
ctx = SessionContext()
for t, i, x in (("a", [3, 2], [0.5, -0.625]), ("b", [2, 0, 3], [0.5, -1.75, -0.375])):
    ctx.register_record_batches(t, [[pa.record_batch({"r": pa.array(range(len(i)), pa.int64()), "i": pa.array(i, pa.int64()), "x": x})]])
q = """WITH u AS (SELECT k, v FROM (SELECT i AS k, x AS v FROM a
                               UNION ALL SELECT k, v FROM (SELECT i AS k, x - 0.5 AS v FROM a
                                                           UNION ALL SELECT i AS k, x AS v FROM a) n1) n2
              UNION ALL SELECT i AS k, x AS v FROM b),
     loss AS (SELECT SUM(v * v) AS l FROM u)
SELECT r, grad(loss, x) AS g FROM b"""
ad.sql(ctx, q).to_arrow_table()   # optimize_unions failed: No field named x

What it takes:

• The data-only union must sit two levels deep. One level deep works.
• Its first branch must compute a column (x - 0.5). With plain x it works.
• If a wrt table is read inside the nested union, it works.

DataFusion itself is fine: the forward query (SELECT l FROM loss) runs. With every step run as SQL (ad._SQL_ABOVE = -1), the query gives the right gradient, [1.0, -3.5, -0.75], as DuckDB does. So the bug is in the Substrait plan the step gets. A guess: #229's renaming of union columns after their qualified names doesn't reach a union that carries no gradient. DataFusion then flattens the nested unions, and the names no longer match.

Found by fuzzing UNION ALL on both engines (a script in my scratch space; I can turn it into a test). In 900 random cases, this failed 6 times.

🤖 Generated with Claude Code

https://claude.ai/code/session_012NE18ox5ivwTc7ZSbHZUGu