[SPARK-37829][SQL] DataFrame.joinWith should return null rows for missing values

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala

@@ -77,6 +77,28 @@ object RowEncoder {
       ClassTag(cls))
   }
 
+  /**
+   * Returns an ExpressionEncoder allowing null top-level rows.
+   * @param exprEnc an ExpressionEncoder[Row].
+   * @return an ExpressionEncoder[Row] whom deserializer supports null values.
+   *
+   * @see SPARK-37829
+   */
+  private[sql] def nullSafe(exprEnc: ExpressionEncoder[Row]): ExpressionEncoder[Row] = {
+    val newDeserializerInput = GetColumnByOrdinal(0, exprEnc.objSerializer.dataType)
+    val newDeserializer: Expression = if (exprEnc.objSerializer.nullable) {
+      If(
+        IsNull(newDeserializerInput),
+        Literal.create(null, exprEnc.objDeserializer.dataType),
+        exprEnc.objDeserializer)
+    } else {
+      exprEnc.objDeserializer
+    }
+    exprEnc.copy(
+      objDeserializer = newDeserializer
+    )
+  }
+
   private def serializerFor(
       inputObject: Expression,
       inputType: DataType): Expression = inputType match {

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -1173,8 +1173,20 @@ class Dataset[T] private[sql](
       joined = resolveSelfJoinCondition(joined)
     }
 
-    implicit val tuple2Encoder: Encoder[(T, U)] =
-      ExpressionEncoder.tuple(this.exprEnc, other.exprEnc)
+    // SPARK-37829: an outer-join requires the null semantics to represent missing keys.
+    // As we might be running on DataFrames, we need a custom encoder that will properly
+    // handle null top-level Rows.
+    def nullSafe[V](exprEnc: ExpressionEncoder[V]): ExpressionEncoder[V] = {
+      if (exprEnc.clsTag.runtimeClass != classOf[Row]) {
+        exprEnc
+      } else {
+        RowEncoder.nullSafe(exprEnc.asInstanceOf[ExpressionEncoder[Row]])
+          .asInstanceOf[ExpressionEncoder[V]]
+      }
+    }
+    implicit val tuple2Encoder: Encoder[(T, U)] = ExpressionEncoder.tuple(
+      nullSafe(this.exprEnc), nullSafe(other.exprEnc)
+    )
 
     val leftResultExpr = {
       if (!this.exprEnc.isSerializedAsStructForTopLevel) {

sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala

@@ -29,6 +29,7 @@ import org.apache.spark.TestUtils.withListener
 import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart}
 import org.apache.spark.sql.catalyst.{FooClassWithEnum, FooEnum, ScroogeLikeExample}
 import org.apache.spark.sql.catalyst.encoders.{OuterScopes, RowEncoder}
+import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
 import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi}
 import org.apache.spark.sql.catalyst.util.sideBySide
 import org.apache.spark.sql.execution.{LogicalRDD, RDDScanExec, SQLExecution}
@@ -2147,6 +2148,49 @@ class DatasetSuite extends QueryTest
       (2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12),
       (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13))
   }
+
+  test("SPARK-37829: DataFrame outer join") {
+    val left = Seq(ClassData("a", 1), ClassData("b", 2)).toDF().as("left")
+    val right = Seq(ClassData("x", 2), ClassData("y", 3)).toDF().as("right")
+    val joined = left.joinWith(right, $"left.b" === $"right.b", "left")
+
+    val leftFieldSchema = StructType(
+      Seq(
+        StructField("a", StringType),
+        StructField("b", IntegerType, nullable = false)
+      )
+    )
+    val rightFieldSchema = StructType(
+      Seq(
+        StructField("a", StringType),
+        StructField("b", IntegerType, nullable = false)
+      )
+    )
+    val expectedSchema = StructType(
+      Seq(
+        StructField(
+          "_1",
+          leftFieldSchema,
+          nullable = false
+        ),
+        // This is a left join, so the right output is nullable:
+        StructField(
+          "_2",
+          rightFieldSchema
+        )
+      )
+    )
+    assert(joined.schema === expectedSchema)
+
+    val result = joined.collect().toSet
+    val expected = Set(
+      new GenericRowWithSchema(Array("a", 1), leftFieldSchema) ->
+        null,
+      new GenericRowWithSchema(Array("b", 2), leftFieldSchema) ->
+        new GenericRowWithSchema(Array("x", 2), rightFieldSchema)
+    )
+    assert(result == expected)
+  }
 }
 
 case class Bar(a: Int)