[SPARK-31071][SQL] Allow annotating non-null fields when encoding Java Beans

### What changes were proposed in this pull request?

When encoding Java Beans to Spark DataFrame, respecting `javax.annotation.Nonnull` and producing non-null fields.

### Why are the changes needed?

When encoding Java Beans to Spark DataFrame, non-primitive types are encoded as nullable fields. Although It works for most cases, it can be an issue under a few situations, e.g. the one described in the JIRA ticket when saving DataFrame to Avro format with non-null field.

We should allow Spark users more flexibility when creating Spark DataFrame from Java Beans. Currently, Spark users cannot create DataFrame with non-nullable fields in the schema from beans with non-nullable properties.

Although it is possible to project top-level columns with SQL expressions like `AssertNotNull` to make it non-null, for nested fields it is more tricky to do it similarly.

### Does this PR introduce any user-facing change?

Yes. After this change, Spark users can use `javax.annotation.Nonnull` to annotate non-null fields in Java Beans when encoding beans to Spark DataFrame.

### How was this patch tested?

Added unit test.

Closes apache#27851 from viirya/SPARK-31071.

Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

Author: viirya

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/JavaTypeInference.scala

@@ -21,6 +21,7 @@ import java.beans.{Introspector, PropertyDescriptor}
 import java.lang.{Iterable => JIterable}
 import java.lang.reflect.Type
 import java.util.{Iterator => JIterator, List => JList, Map => JMap}
+import javax.annotation.Nonnull
 
 import scala.language.existentials
 
@@ -148,7 +149,9 @@ object JavaTypeInference {
         val fields = properties.map { property =>
           val returnType = typeToken.method(property.getReadMethod).getReturnType
           val (dataType, nullable) = inferDataType(returnType, seenTypeSet + other)
-          new StructField(property.getName, dataType, nullable)
+          // The existence of `javax.annotation.Nonnull`, means this field is not nullable.
+          val hasNonNull = property.getReadMethod.isAnnotationPresent(classOf[Nonnull])
+          new StructField(property.getName, dataType, nullable && !hasNonNull)
         }
         (new StructType(fields), true)
     }
@@ -340,10 +343,12 @@ object JavaTypeInference {
           val fieldType = typeToken.method(p.getReadMethod).getReturnType
           val (dataType, nullable) = inferDataType(fieldType)
           val newTypePath = walkedTypePath.recordField(fieldType.getType.getTypeName, fieldName)
+          // The existence of `javax.annotation.Nonnull`, means this field is not nullable.
+          val hasNonNull = p.getReadMethod.isAnnotationPresent(classOf[Nonnull])
           val setter = expressionWithNullSafety(
             deserializerFor(fieldType, addToPath(path, fieldName, dataType, newTypePath),
               newTypePath),
-            nullable = nullable,
+            nullable = nullable && !hasNonNull,
             newTypePath)
           p.getWriteMethod.getName -> setter
         }.toMap
@@ -442,10 +447,13 @@ object JavaTypeInference {
           val fields = properties.map { p =>
             val fieldName = p.getName
             val fieldType = typeToken.method(p.getReadMethod).getReturnType
+            val hasNonNull = p.getReadMethod.isAnnotationPresent(classOf[Nonnull])
             val fieldValue = Invoke(
               inputObject,
               p.getReadMethod.getName,
-              inferExternalType(fieldType.getRawType))
+              inferExternalType(fieldType.getRawType),
+              propagateNull = !hasNonNull,
+              returnNullable = !hasNonNull)
             (fieldName, serializerFor(fieldValue, fieldType))
           }
           createSerializerForObject(inputObject, fields)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SerializerBuildHelper.scala

@@ -187,8 +187,12 @@ object SerializerBuildHelper {
     val nonNullOutput = CreateNamedStruct(fields.flatMap { case(fieldName, fieldExpr) =>
       argumentsForFieldSerializer(fieldName, fieldExpr)
     })
-    val nullOutput = expressions.Literal.create(null, nonNullOutput.dataType)
-    expressions.If(IsNull(inputObject), nullOutput, nonNullOutput)
+    if (inputObject.nullable) {
+      val nullOutput = expressions.Literal.create(null, nonNullOutput.dataType)
+      expressions.If(IsNull(inputObject), nullOutput, nonNullOutput)
+    } else {
+      nonNullOutput
+    }
   }
 
   def createSerializerForUserDefinedType(

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

@@ -203,6 +203,7 @@ case class ExpressionEncoder[T](
       }
       nullSafeSerializer match {
         case If(_: IsNull, _, s: CreateNamedStruct) => s
+        case s: CreateNamedStruct => s
         case _ =>
           throw new RuntimeException(s"class $clsName has unexpected serializer: $objSerializer")
       }

sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java

@@ -24,6 +24,7 @@
 import java.time.Instant;
 import java.time.LocalDate;
 import java.util.*;
+import javax.annotation.Nonnull;
 
 import org.apache.spark.sql.streaming.GroupStateTimeout;
 import org.apache.spark.sql.streaming.OutputMode;
@@ -823,6 +824,75 @@ public int hashCode() {
     }
   }
 
+  public static class NestedSmallBeanWithNonNullField implements Serializable {
+    private SmallBean nonNull_f;
+    private SmallBean nullable_f;
+    private Map<String, SmallBean> childMap;
+
+    @Nonnull
+    public SmallBean getNonNull_f() {
+      return nonNull_f;
+    }
+
+    public void setNonNull_f(SmallBean f) {
+      this.nonNull_f = f;
+    }
+
+    public SmallBean getNullable_f() {
+      return nullable_f;
+    }
+
+    public void setNullable_f(SmallBean f) {
+      this.nullable_f = f;
+    }
+
+    @Nonnull
+    public Map<String, SmallBean> getChildMap() { return childMap; }
+    public void setChildMap(Map<String, SmallBean> childMap) {
+      this.childMap = childMap;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      NestedSmallBeanWithNonNullField that = (NestedSmallBeanWithNonNullField) o;
+      return Objects.equal(nullable_f, that.nullable_f) &&
+        Objects.equal(nonNull_f, that.nonNull_f) && Objects.equal(childMap, that.childMap);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(nullable_f, nonNull_f, childMap);
+    }
+  }
+
+  public static class NestedSmallBean2 implements Serializable {
+    private NestedSmallBeanWithNonNullField f;
+
+    @Nonnull
+    public NestedSmallBeanWithNonNullField getF() {
+      return f;
+    }
+
+    public void setF(NestedSmallBeanWithNonNullField f) {
+      this.f = f;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      NestedSmallBean2 that = (NestedSmallBean2) o;
+      return Objects.equal(f, that.f);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(f);
+    }
+  }
+
   @Rule
   public transient ExpectedException nullabilityCheck = ExpectedException.none();
 
@@ -1504,6 +1574,54 @@ public void testSerializeNull() {
     Assert.assertEquals(beans, ds2.collectAsList());
   }
 
+  @Test
+  public void testNonNullField() {
+    NestedSmallBeanWithNonNullField bean1 = new NestedSmallBeanWithNonNullField();
+    SmallBean smallBean = new SmallBean();
+    bean1.setNonNull_f(smallBean);
+    Map<String, SmallBean> map = new HashMap<>();
+    bean1.setChildMap(map);
+
+    Encoder<NestedSmallBeanWithNonNullField> encoder1 =
+            Encoders.bean(NestedSmallBeanWithNonNullField.class);
+    List<NestedSmallBeanWithNonNullField> beans1 = Arrays.asList(bean1);
+    Dataset<NestedSmallBeanWithNonNullField> ds1 = spark.createDataset(beans1, encoder1);
+
+    StructType schema = ds1.schema();
+    Assert.assertFalse(schema.apply("nonNull_f").nullable());
+    Assert.assertTrue(schema.apply("nullable_f").nullable());
+    Assert.assertFalse(schema.apply("childMap").nullable());
+
+    Assert.assertEquals(beans1, ds1.collectAsList());
+    Dataset<NestedSmallBeanWithNonNullField> ds2 = ds1.map(
+      (MapFunction<NestedSmallBeanWithNonNullField, NestedSmallBeanWithNonNullField>) b -> b,
+      encoder1);
+    Assert.assertEquals(beans1, ds2.collectAsList());
+
+    // Nonnull nested fields
+    NestedSmallBean2 bean2 = new NestedSmallBean2();
+    bean2.setF(bean1);
+
+    Encoder<NestedSmallBean2> encoder2 =
+            Encoders.bean(NestedSmallBean2.class);
+    List<NestedSmallBean2> beans2 = Arrays.asList(bean2);
+    Dataset<NestedSmallBean2> ds3 = spark.createDataset(beans2, encoder2);
+
+    StructType nestedSchema = (StructType) ds3.schema()
+      .fields()[ds3.schema().fieldIndex("f")]
+      .dataType();
+    Assert.assertFalse(nestedSchema.apply("nonNull_f").nullable());
+    Assert.assertTrue(nestedSchema.apply("nullable_f").nullable());
+    Assert.assertFalse(nestedSchema.apply("childMap").nullable());
+
+    Assert.assertEquals(beans2, ds3.collectAsList());
+
+    Dataset<NestedSmallBean2> ds4 = ds3.map(
+      (MapFunction<NestedSmallBean2, NestedSmallBean2>) b -> b,
+      encoder2);
+    Assert.assertEquals(beans2, ds4.collectAsList());
+  }
+
   @Test
   public void testSpecificLists() {
     SpecificListsBean bean = new SpecificListsBean();