[branch-54] fix: preserve null_aware on logical JoinNode proto round-trip (backport apache#22104) (apache#22785)

## Which issue does this PR close?

- Backport of apache#22104 to `branch-54` (for 54.1.0, tracked in apache#22547).

This PR:
- Backports apache#22104 to the `branch-54` line so the `null_aware` proto
  round-trip fix ships in 54.1.0, as requested in

apache#22065 (comment)

Clean cherry-pick; `datafusion-proto` builds and both round-trip
regression tests pass on `branch-54`.

Author: mithuncy

datafusion/proto/proto/datafusion.proto

@@ -252,6 +252,7 @@ message JoinNode {
   repeated LogicalExprNode right_join_key = 6;
   datafusion_common.NullEquality null_equality = 7;
   LogicalExprNode filter = 8;
+  bool null_aware = 9;
 }
 
 message DistinctNode {

datafusion/proto/src/generated/pbjson.rs

@@ -64,9 +64,8 @@ use datafusion_expr::{
     Statement, WindowUDF, dml,
     logical_plan::{
         Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView,
-        DdlStatement, Distinct, EmptyRelation, Extension, Join, JoinConstraint, Prepare,
-        Projection, Repartition, Sort, SubqueryAlias, TableScan, Values, Window,
-        builder::project,
+        DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection,
+        Repartition, Sort, SubqueryAlias, TableScan, Values, Window, builder::project,
     },
 };
 
@@ -850,6 +849,13 @@ impl AsLogicalPlan for LogicalPlanNode {
                     from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?;
                 let right_keys: Vec<Expr> =
                     from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?;
+                if left_keys.len() != right_keys.len() {
+                    return Err(proto_error(format!(
+                        "Received a JoinNode message with left_join_key and right_join_key of different lengths: {} and {}",
+                        left_keys.len(),
+                        right_keys.len()
+                    )));
+                }
                 let join_type =
                     protobuf::JoinType::try_from(join.join_type).map_err(|_| {
                         proto_error(format!(
@@ -866,44 +872,39 @@ impl AsLogicalPlan for LogicalPlanNode {
                         join.join_constraint
                     ))
                 })?;
+                let null_equality = protobuf::NullEquality::try_from(join.null_equality)
+                    .map_err(|_| {
+                        proto_error(format!(
+                            "Received a JoinNode message with unknown NullEquality {}",
+                            join.null_equality
+                        ))
+                    })?;
                 let filter: Option<Expr> = join
                     .filter
                     .as_ref()
                     .map(|expr| from_proto::parse_expr(expr, ctx, extension_codec))
                     .map_or(Ok(None), |v| v.map(Some))?;
-
-                let builder = LogicalPlanBuilder::from(into_logical_plan!(
-                    join.left,
-                    ctx,
-                    extension_codec
-                )?);
-                let builder = match join_constraint.into() {
-                    JoinConstraint::On => builder.join_with_expr_keys(
-                        into_logical_plan!(join.right, ctx, extension_codec)?,
-                        join_type.into(),
-                        (left_keys, right_keys),
-                        filter,
-                    )?,
-                    JoinConstraint::Using => {
-                        // The equijoin keys in using-join must be column.
-                        let using_keys = left_keys
-                            .into_iter()
-                            .map(|key| {
-                                key.try_as_col().cloned()
-                                    .ok_or_else(|| internal_datafusion_err!(
-                                        "Using join keys must be column references, got: {key:?}"
-                                    ))
-                            })
-                            .collect::<Result<Vec<_>, _>>()?;
-                        builder.join_using(
-                            into_logical_plan!(join.right, ctx, extension_codec)?,
-                            join_type.into(),
-                            using_keys,
-                        )?
-                    }
-                };
-
-                builder.build()
+                let left = into_logical_plan!(join.left, ctx, extension_codec)?;
+                let right = into_logical_plan!(join.right, ctx, extension_codec)?;
+                let on: Vec<(Expr, Expr)> =
+                    left_keys.into_iter().zip(right_keys).collect();
+
+                // Construct the Join directly instead of going through
+                // LogicalPlanBuilder. The builder methods hardcode
+                // `null_equality` and `null_aware`, so a round trip through
+                // them silently loses both fields. Both sides of the round
+                // trip should already have validated keys, so we don't need
+                // the builder's normalization / equijoin-pair checks.
+                Ok(LogicalPlan::Join(Join::try_new(
+                    Arc::new(left),
+                    Arc::new(right),
+                    on,
+                    filter,
+                    join_type.into(),
+                    join_constraint.into(),
+                    null_equality.into(),
+                    join.null_aware,
+                )?))
             }
             LogicalPlanType::Union(union) => {
                 assert_or_internal_err!(
@@ -1492,7 +1493,9 @@ impl AsLogicalPlan for LogicalPlanNode {
                 join_type,
                 join_constraint,
                 null_equality,
-                ..
+                null_aware,
+                // Not encoded; recomputed by `Join::try_new` on decode.
+                schema: _,
             }) => {
                 let left: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
                     left.as_ref(),
@@ -1533,6 +1536,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                             right_join_key,
                             null_equality: null_equality.into(),
                             filter,
+                            null_aware: *null_aware,
                         },
                     ))),
                 })

datafusion/proto/src/generated/prost.rs

@@ -3173,3 +3173,88 @@ async fn roundtrip_empty_table_scan_with_projection() -> Result<()> {
     );
     Ok(())
 }
+
+// Regression test for https://github.com/apache/datafusion/issues/22065:
+// the decoder must preserve `null_aware = true` (NOT IN semantics)
+// across a to_proto -> from_proto round trip. `null_equality` is at
+// its default (`NullEqualsNothing`).
+#[tokio::test]
+async fn roundtrip_join_null_aware() -> Result<()> {
+    use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
+    use datafusion_expr::JoinType;
+
+    let ctx = SessionContext::new();
+    let sql = "
+        SELECT id
+        FROM (VALUES (1), (2), (3)) AS t1(id)
+        WHERE id NOT IN (
+            SELECT bad_id
+            FROM (VALUES (CAST(1 AS INT)), (CAST(NULL AS INT))) AS excludes(bad_id)
+        )
+    ";
+
+    let df = ctx.sql(sql).await?;
+    let plan = ctx.state().optimize(df.logical_plan())?;
+
+    let mut found_null_aware = false;
+    plan.apply(|n| {
+        if let LogicalPlan::Join(j) = n
+            && j.join_type == JoinType::LeftAnti
+            && j.null_aware
+        {
+            found_null_aware = true;
+        }
+        Ok(TreeNodeRecursion::Continue)
+    })?;
+    assert!(found_null_aware);
+
+    let bytes = logical_plan_to_bytes(&plan)?;
+    let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?;
+    assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}"));
+
+    Ok(())
+}
+
+// Regression test for `null_equality` round-trip (related to #22065):
+// the decoder must preserve a non-default `null_equality`
+// (`NullEqualsNull`) across a to_proto -> from_proto round trip.
+// `null_aware` is at its default (`false`).
+#[tokio::test]
+async fn roundtrip_join_null_equality() -> Result<()> {
+    use datafusion_common::NullEquality;
+    use datafusion_expr::JoinType;
+    use datafusion_expr::logical_plan::{Join, JoinConstraint};
+
+    let ctx = SessionContext::new();
+
+    let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
+    let right_schema =
+        Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)]));
+    ctx.register_table(
+        "t1",
+        Arc::new(datafusion::datasource::empty::EmptyTable::new(left_schema)),
+    )?;
+    ctx.register_table(
+        "t2",
+        Arc::new(datafusion::datasource::empty::EmptyTable::new(right_schema)),
+    )?;
+    let left = ctx.table("t1").await?.into_optimized_plan()?;
+    let right = ctx.table("t2").await?.into_optimized_plan()?;
+
+    let join = LogicalPlan::Join(Join::try_new(
+        Arc::new(left),
+        Arc::new(right),
+        vec![(col("t1.a"), col("t2.b"))],
+        None,
+        JoinType::Inner,
+        JoinConstraint::On,
+        NullEquality::NullEqualsNull,
+        false,
+    )?);
+
+    let bytes = logical_plan_to_bytes(&join)?;
+    let rt = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?;
+    assert_eq!(format!("{join:?}"), format!("{rt:?}"));
+
+    Ok(())
+}