Add plan assert as well

Author: hareshkh

datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs

@@ -101,7 +101,9 @@ fn split_eq_and_noneq_join_predicate_with_nulls_equality(
             }) => match (left.as_ref(), right.as_ref()) {
                 (Expr::Column(l), Expr::Column(r)) => match op {
                     Operator::Eq => eq_keys.push((l.clone(), r.clone())),
-                    Operator::IsNotDistinctFrom => indistinct_keys.push((l.clone(), r.clone())),
+                    Operator::IsNotDistinctFrom => {
+                        indistinct_keys.push((l.clone(), r.clone()))
+                    }
                     _ => unreachable!(),
                 },
                 _ => accum_filters.push(expr.clone()),
@@ -110,26 +112,26 @@ fn split_eq_and_noneq_join_predicate_with_nulls_equality(
         }
     }
 
-    let (join_keys, null_equality) = match (eq_keys.is_empty(), indistinct_keys.is_empty())
-    {
-        // Mixed: use eq_keys as equijoin keys, demote indistinct keys to filter
-        (false, false) => {
-            for (l, r) in &indistinct_keys {
-                accum_filters.push(Expr::BinaryExpr(BinaryExpr {
-                    left: Box::new(Expr::Column(l.clone())),
-                    op: Operator::IsNotDistinctFrom,
-                    right: Box::new(Expr::Column(r.clone())),
-                }));
+    let (join_keys, null_equality) =
+        match (eq_keys.is_empty(), indistinct_keys.is_empty()) {
+            // Mixed: use eq_keys as equijoin keys, demote indistinct keys to filter
+            (false, false) => {
+                for (l, r) in &indistinct_keys {
+                    accum_filters.push(Expr::BinaryExpr(BinaryExpr {
+                        left: Box::new(Expr::Column(l.clone())),
+                        op: Operator::IsNotDistinctFrom,
+                        right: Box::new(Expr::Column(r.clone())),
+                    }));
+                }
+                (eq_keys, NullEquality::NullEqualsNothing)
             }
-            (eq_keys, NullEquality::NullEqualsNothing)
-        }
-        // Only eq keys
-        (false, true) => (eq_keys, NullEquality::NullEqualsNothing),
-        // Only indistinct keys
-        (true, false) => (indistinct_keys, NullEquality::NullEqualsNull),
-        // No keys at all
-        (true, true) => (vec![], NullEquality::NullEqualsNothing),
-    };
+            // Only eq keys
+            (false, true) => (eq_keys, NullEquality::NullEqualsNothing),
+            // Only indistinct keys
+            (true, false) => (indistinct_keys, NullEquality::NullEqualsNull),
+            // No keys at all
+            (true, true) => (vec![], NullEquality::NullEqualsNothing),
+        };
 
     let join_filter = accum_filters.into_iter().reduce(Expr::and);
     (join_keys, null_equality, join_filter)
@@ -224,7 +226,8 @@ mod tests {
         assert_eq!(null_eq, NullEquality::NullEqualsNothing);
 
         // The IsNotDistinctFrom predicate should be demoted to the filter.
-        let filter = filter.expect("filter should contain the demoted indistinct predicate");
+        let filter =
+            filter.expect("filter should contain the demoted indistinct predicate");
         match &filter {
             Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
                 assert_eq!(*op, Operator::IsNotDistinctFrom);

datafusion/substrait/tests/cases/consumer_integration.rs

@@ -763,37 +763,38 @@ mod tests {
         Ok(())
     }
 
-    /// Regression test: a Substrait join expression containing both `equal` and
-    /// `is_not_distinct_from` (as Spark can produce) must preserve the
-    /// null-safe semantics of `IS NOT DISTINCT FROM` by demoting it to the
-    /// join filter when mixed with regular equality keys.
-    ///
-    /// The plan is loaded from a JSON-encoded Substrait protobuf to exercise
-    /// the full consumer path (`from_substrait_plan` → `from_join_rel`).
+    /// Regression: a Substrait join with both `equal` and `is_not_distinct_from`
+    /// must demote `IS NOT DISTINCT FROM` to the join filter (matching the SQL
+    /// planner behavior tested in `join_is_not_distinct_from.slt:179-205`).
     #[tokio::test]
     async fn test_mixed_join_equal_and_indistinct_from_substrait_plan() -> Result<()> {
+        let plan_str =
+            test_plan_to_string("mixed_join_equal_and_indistinct.json").await?;
+        // Eq becomes the equijoin key; IS NOT DISTINCT FROM is demoted to filter.
+        assert_snapshot!(
+            plan_str,
+            @r#"
+        Projection: left.id, left.val, left.comment, right.id AS id0, right.val AS val0, right.comment AS comment0
+          Inner Join: left.id = right.id Filter: left.val IS NOT DISTINCT FROM right.val
+            SubqueryAlias: left
+              Values: (Utf8("1"), Utf8("a"), Utf8("c1")), (Utf8("2"), Utf8("b"), Utf8("c2")), (Utf8("3"), Utf8(NULL), Utf8("c3")), (Utf8("4"), Utf8(NULL), Utf8("c4")), (Utf8("5"), Utf8("e"), Utf8("c5"))...
+            SubqueryAlias: right
+              Values: (Utf8("1"), Utf8("a"), Utf8("c1")), (Utf8("2"), Utf8("b"), Utf8("c2")), (Utf8("3"), Utf8(NULL), Utf8("c3")), (Utf8("4"), Utf8(NULL), Utf8("c4")), (Utf8("5"), Utf8("e"), Utf8("c5"))...
+        "#
+        );
+
+        // Also execute to verify NULL=NULL rows (ids 3,4) are preserved.
         let path = "tests/testdata/test_plans/mixed_join_equal_and_indistinct.json";
         let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
             File::open(path).expect("file not found"),
         ))
         .expect("failed to parse json");
-
         let ctx = SessionContext::new();
         let plan = from_substrait_plan(&ctx.state(), &proto).await?;
-
-        // Execute and count rows.
-        // Both tables have 6 identical rows; rows 3 and 4 have val=NULL.
-        // With correct handling, IS NOT DISTINCT FROM is demoted to the join
-        // filter, so NULL=NULL matches and all 6 rows appear in the output.
         let df = ctx.execute_logical_plan(plan).await?;
         let results = df.collect().await?;
         let total_rows: usize = results.iter().map(|b| b.num_rows()).sum();
-
-        assert_eq!(
-            total_rows, 6,
-            "Expected 6 rows (including NULL=NULL matches via IS NOT DISTINCT FROM), \
-             got {total_rows}. Mixed equal/is_not_distinct_from lost null-safe semantics."
-        );
+        assert_eq!(total_rows, 6);
 
         Ok(())
     }