feat(ddx-ad): differentiate through UNION ALL in reverse mode

crates/ddx-ad/src/contraction.rs

@@ -871,6 +871,15 @@ fn slice(region: &Region, leaf: &Leaf) -> Result<Option<Region>> {
                 }
             }
             Def::Const => Def::Const,
+            Def::Branch => Def::Branch,
+            Def::Union(arms) => match arms
+                .iter()
+                .map(|&(b, c)| Ok((local(b)?, local(c)?)))
+                .collect::<Result<Vec<_>>>()
+            {
+                Ok(arms) => Def::Union(arms),
+                Err(_) => return Ok(None),
+            },
         };
         defs.push(def);
     }

crates/ddx-ad/src/emit.rs

@@ -5,8 +5,8 @@
 //! Writing the plans of a backward program.
 //!
 //! Every step ddx emits is an ordinary Substrait plan, with nothing in it only
-//! ddx understands, built from six relations: read, project, filter, join,
-//! aggregate and union. The builders here produce them in the shapes
+//! ddx understands, built from seven relations: read, project, filter, join,
+//! aggregate, union and fetch. The builders here produce them in the shapes
 //! DataFusion's own producer does, since that is the consumer they are tested
 //! against.
 //!
@@ -25,14 +25,15 @@
 use substrait::proto::aggregate_function::AggregationInvocation;
 use substrait::proto::aggregate_rel::{Grouping, Measure};
 use substrait::proto::expression::MaskExpression;
+use substrait::proto::fetch_rel::CountMode;
 use substrait::proto::join_rel::JoinType;
 use substrait::proto::plan_rel::RelType as PlanRelType;
 use substrait::proto::read_rel::{NamedTable, ReadType};
 use substrait::proto::rel::RelType;
 use substrait::proto::rel_common::{Emit, EmitKind};
 use substrait::proto::set_rel::SetOp;
 use substrait::proto::{
-    r#type, AggregateFunction, AggregateRel, Expression, FilterRel, FunctionArgument, JoinRel,
+    r#type, AggregateFunction, AggregateRel, Expression, FetchRel, FilterRel, FunctionArgument, JoinRel,
     NamedStruct, Plan, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, RelRoot, SetRel,
 };
 
@@ -201,6 +202,20 @@ pub fn union_all(inputs: Vec<Rel>) -> Rel {
     }
 }
 
+/// None of `input`'s rows: its columns, typed as the engine types them,
+/// with nothing in them. Joined to with a FULL JOIN, it gives another
+/// relation's rows NULL in those columns.
+pub fn no_rows(input: Rel) -> Rel {
+    Rel {
+        rel_type: Some(RelType::Fetch(Box::new(FetchRel {
+            input: Some(Box::new(input)),
+            #[allow(deprecated)]
+            count_mode: Some(CountMode::Count(0)),
+            ..Default::default()
+        }))),
+    }
+}
+
 /// A plan whose root is `root`, with output columns named `names`.
 pub fn plan(mut root: Rel, names: Vec<String>, ext: &Extensions) -> Plan {
     condition_joins(&mut root);

crates/ddx-ad/src/forward.rs

@@ -53,9 +53,10 @@ use substrait::proto::plan_rel::RelType as PlanRelType;
 use substrait::proto::read_rel::ReadType;
 use substrait::proto::rel::RelType;
 use substrait::proto::rel_common::EmitKind;
+use substrait::proto::set_rel::SetOp;
 use substrait::proto::{
     AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, Expression, FetchRel,
-    JoinRel, NamedStruct, Plan, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, SortRel,
+    JoinRel, NamedStruct, Plan, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, SetRel, SortRel,
 };
 
 use crate::elementwise::depends;
@@ -158,6 +159,21 @@ pub enum Def {
     },
     /// A column of a constant input.
     Const,
+    /// TRUE in the rows one input of a `UNION ALL` gave, NULL in the rows
+    /// the others gave. See [`Def::Union`].
+    Branch,
+    /// A column of a `UNION ALL`: in the rows each input gave, that input's
+    /// column. Each pair is an input's [`Def::Branch`] column and its column.
+    ///
+    /// A union is rebuilt as one relation that holds every input's columns:
+    /// each input's rows carry its own columns and NULL in every other
+    /// input's. So the row-local rules apply above and beneath it, and its
+    /// transpose is a select: a row's cotangent goes to the column of the
+    /// input that gave the row, and NULL, no cotangent, to the others'.
+    /// Everything beneath an input is computed from its own columns, so its
+    /// cotangents are NULL in the rows other inputs gave, and those rows
+    /// send it no gradient.
+    Union(Vec<(usize, usize)>),
 }
 
 /// A rebuilt, recomputed region of row-local work.
@@ -214,6 +230,12 @@ impl Region {
                     offset: offset + shift,
                 },
                 Def::Const => Def::Const,
+                Def::Branch => Def::Branch,
+                Def::Union(arms) => Def::Union(
+                    arms.into_iter()
+                        .map(|(b, c)| (b + shift, c + shift))
+                        .collect(),
+                ),
             });
         }
         self.varied.extend(other.varied);
@@ -1003,6 +1025,10 @@ impl Builder<'_> {
                 let (s, direct) = self.cross(c)?;
                 (s, direct, c.common.as_ref())
             }
+            RelType::Set(u) => {
+                let (s, direct) = self.union_all(u)?;
+                (s, direct, u.common.as_ref())
+            }
             RelType::Aggregate(a) => {
                 let n = self.read_saved(a, below.take())?;
                 let width = self.saved[n].outputs.len();
@@ -1305,6 +1331,90 @@ impl Builder<'_> {
         Ok((s, direct))
     }
 
+    /// A `UNION ALL`, rebuilt as one relation that holds every input's
+    /// columns (see [`Def::Union`]): each input's columns, then its
+    /// [`Def::Branch`]; after them all, the union's columns.
+    fn union_all(&mut self, u: &SetRel) -> Result<(Region, Vec<usize>)> {
+        if SetOp::try_from(u.op) != Ok(SetOp::UnionAll) {
+            return Err(AdError::NotImplemented(
+                "a set operation other than UNION ALL on the path from a wrt table to the \
+                 output"
+                    .into(),
+            ));
+        }
+        let mut inputs = Vec::with_capacity(u.inputs.len());
+        for r in &u.inputs {
+            let mut s = self.lower(r)?;
+            s.rel = emit::project(s.rel, vec![lit_bool(true)]);
+            s.push(Def::Branch, false);
+            inputs.push(s);
+        }
+        let rels: Vec<Rel> = inputs.iter().map(|s| s.rel.clone()).collect();
+        let mut inputs = inputs.into_iter();
+        let mut s = inputs
+            .next()
+            .ok_or_else(|| AdError::InvalidPlan("a set operation with no inputs".into()))?;
+        // Each input's first column here, its outputs, and its branch column.
+        let mut blocks = vec![(0, s.outputs.clone(), s.width() - 1)];
+        for r in inputs {
+            let outputs = r.outputs.clone();
+            let w = r.width();
+            let shift = s.append(r)?;
+            blocks.push((shift, outputs, shift + w - 1));
+        }
+        let width = blocks[0].1.len();
+        if blocks.iter().any(|(_, outputs, _)| outputs.len() != width) {
+            return Err(AdError::InvalidPlan(
+                "a UNION ALL whose inputs have different numbers of columns".into(),
+            ));
+        }
+        // Each input's rows, with every other input's columns NULL: FULL
+        // JOINed to none of the others' rows, so the engine types the NULLs
+        // as it types the columns.
+        //
+        // DataFusion names a union's columns after its first input's,
+        // without their tables, and suffixes a repeated name: a second `x`
+        // is `x_1`, though a table may already have an `x_1`, and the union
+        // is refused. So each input's columns are renamed first, each after
+        // its own table and column, which no other has.
+        let w = s.width() + width;
+        let all: Vec<usize> = (0..w).collect();
+        let mut rows = Vec::with_capacity(rels.len());
+        for (k, (offset, outputs, _)) in blocks.iter().enumerate() {
+            let row = rels
+                .iter()
+                .enumerate()
+                .map(|(j, r)| if j == k { r.clone() } else { emit::no_rows(r.clone()) })
+                .reduce(|l, r| emit::join(l, r, lit_bool(true), JoinType::Outer))
+                .expect("at least one input");
+            let row = emit::project(row, outputs.iter().map(|&c| field(offset + c)).collect());
+            rows.push(match rels.len() {
+                1 => row,
+                _ => rename_in_place(row, w, &all),
+            });
+        }
+        s.rel = match rows.len() {
+            1 => rows.pop().expect("one input"),
+            _ => emit::union_all(rows),
+        };
+        let base = s.width();
+        for c in 0..width {
+            let arms: Vec<(usize, usize)> = blocks
+                .iter()
+                .map(|(offset, outputs, branch)| (*branch, offset + outputs[c]))
+                .collect();
+            let varied = arms.iter().any(|&(_, a)| s.varied[a]);
+            s.push(Def::Union(arms), varied);
+        }
+        // The union has every input's rows: one row from each is several.
+        if blocks.len() > 1 {
+            for slot in &mut s.slots {
+                slot.at_most_one_row = false;
+            }
+        }
+        Ok((s, (base..base + width).collect()))
+    }
+
     /// The saved relation for aggregate `a`, read on first sight.
     ///
     /// A query that reads a CTE twice gets the CTE's plan twice: Substrait

crates/ddx-ad/src/transpose.rs

@@ -13,6 +13,7 @@
 //! | **map** | a projected expression `y = f(x₁, x₂, …)` | `x̄ᵢ += ȳ · ∂f/∂xᵢ`, row by row ([`crate::elementwise::Elementwise`]) |
 //! | **select** | `WHERE`, a join condition, a semi-join, a filter on a rank | the cotangent stays on the rows that were kept |
 //! | **broadcast** | a join, which pairs each row with every row it matches | sum the cotangent back over the rows each input row was copied to |
+//! | **concatenate** | `UNION ALL` | each row's cotangent to the input that gave the row |
 //! | **reduce** | a grouped `SUM` | broadcast the group's cotangent to every row that was summed |
 //! | | `AVG` | the same, divided by the group's count |
 //! | | `MAX`, `MIN` | the group's cotangent to the rows that attain it, shared evenly at a tie (`jax.grad`'s convention for `jnp.max`) |
@@ -313,6 +314,8 @@ impl<'a> Transposer<'a> {
                 }
             }
         }
+        // A UNION ALL's inputs can give rows with the same dims.
+        keys.extend((0..beneath.defs.len()).filter(|&c| matches!(beneath.defs[c], Def::Branch)));
         Ok(keys)
     }
 
@@ -650,7 +653,19 @@ impl<'a> Transposer<'a> {
                     rel = project(rel, std::mem::take(&mut batch));
                     self.gather(beneath, *offset, function, &rel, here)?;
                 }
-                Def::Const => {
+                Def::Union(arms) => {
+                    // Each input's column gets the cotangent in the rows it
+                    // gave, and NULL, none, in the others'. An input whose
+                    // column reads no wrt table gets none, as a map's
+                    // constant operand does.
+                    for &(branch, col) in arms.iter().filter(|&&(_, col)| region.varied[col]) {
+                        pending
+                            .entry(col)
+                            .or_default()
+                            .push(if_then(vec![(field(branch), field(here))], null_f64()));
+                    }
+                }
+                Def::Const | Def::Branch => {
                     return Err(AdError::Internal(format!(
                         "gradient reached column {c}, a {:?}, with no refusal recorded",
                         region.defs[c]
@@ -985,6 +1000,8 @@ fn jitters(f: &Forward, region: &Region, col: usize) -> bool {
                 }
             }
             Def::Window { .. } => return true,
+            Def::Branch => {}
+            Def::Union(arms) => stack.extend(arms.iter().map(|&(_, col)| col)),
         }
     }
     false