[SPARK-47927][SQL][FOLLOWUP] fix ScalaUDF output nullability (alternative fix)

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -1697,7 +1697,7 @@ class SparkConnectPlanner(
         inputEncoders = udfPacket.inputEncoders.map(e => Try(ExpressionEncoder(e)).toOption),
         outputEncoder = Option(ExpressionEncoder(udfPacket.outputEncoder)),
         udfName = Option(fun.getFunctionName),
-        nullable = udf.getNullable,
+        outputNullable = udf.getNullable,
         udfDeterministic = fun.getDeterministic)
     }
   }
@@ -2071,7 +2071,7 @@ class SparkConnectPlanner(
       inputEncoders = f.inputEncoders,
       outputEncoder = f.outputEncoder,
       udfName = f.name,
-      nullable = f.nullable,
+      outputNullable = f.nullable,
       udfDeterministic = f.deterministic)
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala

@@ -52,10 +52,14 @@ case class ScalaUDF(
     inputEncoders: Seq[Option[ExpressionEncoder[_]]] = Nil,
     outputEncoder: Option[ExpressionEncoder[_]] = None,
     udfName: Option[String] = None,
-    nullable: Boolean = true,
+    outputNullable: Boolean = true,
     udfDeterministic: Boolean = true)
   extends Expression with NonSQLExpression with UserDefinedExpression {
 
+  // The Rule HandleNullInputsForUDF makes the UDF null-propagatable for primitive nullable inputs.
+  override lazy val nullable: Boolean = outputNullable || inputPrimitives.zip(children)
+    .exists {case (isPrimitive, child) => isPrimitive && child.nullable}
+
   override lazy val deterministic: Boolean = udfDeterministic && children.forall(_.deterministic)
 
   // `ScalaUDF` uses `ExpressionEncoder` to convert the function result to Catalyst internal format.

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala

@@ -1796,4 +1796,15 @@ class AnalysisSuite extends AnalysisTest with Matchers {
     assert(refs.head.resolved)
     assert(refs.head.isStreaming)
   }
+
+  test("SPARK-47927: ScalaUDF output nullability") {
+    val udf = ScalaUDF(
+      function = (i: Int) => i + 1,
+      dataType = IntegerType,
+      children = $"a" :: Nil,
+      outputNullable = false,
+      inputEncoders = Seq(Some(ExpressionEncoder[Int]().resolveAndBind())))
+    val plan = testRelation.select(udf.as("u")).select($"u").analyze
+    assert(plan.output.head.nullable)
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/expressions/UserDefinedFunction.scala

@@ -108,7 +108,7 @@ private[spark] case class SparkUserDefinedFunction(
       inputEncoders,
       outputEncoder,
       udfName = name,
-      nullable = nullable,
+      outputNullable = nullable,
       udfDeterministic = deterministic)
   }