test: a property test for identity; a one-row join of constant data

crates/ddx-ad/src/forward.rs

@@ -1746,7 +1746,10 @@ pub(crate) fn read_outputs(r: &ReadRel) -> Result<Vec<usize>> {
 
 /// Does the plan guarantee `rel` has at most one row? True for an aggregate
 /// with no grouping, a one-row `VALUES`, and anything that only filters,
-/// projects, sorts, limits or cross-joins such relations. False when unsure.
+/// projects, sorts, limits or joins such relations (DataFusion writes a
+/// `CROSS JOIN` as a join on `true`). A full outer join of two can give
+/// two rows; a semi- or anti-join gives at most its kept side's. False when
+/// unsure.
 fn at_most_one_row(rel: &Rel) -> bool {
     let Some(kind) = &rel.rel_type else {
         return false;
@@ -1759,6 +1762,14 @@ fn at_most_one_row(rel: &Rel) -> bool {
         RelType::Sort(s) => below(&s.input),
         RelType::Fetch(f) => below(&f.input),
         RelType::Cross(c) => below(&c.left) && below(&c.right),
+        RelType::Join(j) => match JoinType::try_from(j.r#type) {
+            Ok(JoinType::Inner | JoinType::Left | JoinType::Right) => {
+                below(&j.left) && below(&j.right)
+            }
+            Ok(JoinType::LeftSemi | JoinType::LeftAnti) => below(&j.left),
+            Ok(JoinType::RightSemi | JoinType::RightAnti) => below(&j.right),
+            _ => false,
+        },
         RelType::Read(r) => match &r.read_type {
             Some(ReadType::VirtualTable(v)) => {
                 // Older producers fill the deprecated `values`.

crates/ddx-datafusion/tests/grad.rs

@@ -142,6 +142,24 @@ async fn a_loss_may_divide_by_a_one_row_constant() {
     .await;
 }
 
+#[tokio::test]
+async fn a_loss_of_one_row_aggregates_joined_is_one_row_when_one_is_constant() {
+    // Found by the identity property test (identity.rs). DataFusion writes
+    // a CROSS JOIN as a join on `true`, and the two aggregates that read no
+    // wrt table are copied in as one constant input, whose join ddx did not
+    // see is one row: the loss was refused as not scalar.
+    check_gradients(
+        &ctx(),
+        "SELECT a0.l + a1.l + a2.l AS loss \
+         FROM (SELECT SUM(w.val * x.v) AS l FROM w JOIN x ON w.i = x.i) a0 \
+         CROSS JOIN (SELECT SUM(val) AS l FROM w) a1 \
+         CROSS JOIN (SELECT SUM(val * val) AS l FROM b) a2",
+        &[w(), x(), b()],
+        &[wrt("b", "val")],
+    )
+    .await;
+}
+
 #[tokio::test]
 async fn a_wrt_table_the_loss_does_not_depend_on_gets_zeros() {
     let grads = check_gradients(