From e21eab33e04d2f0dcb1791e98bffa01eda986c6d Mon Sep 17 00:00:00 2001
From: Kent Yao <yao@apache.org>
Date: Fri, 15 Mar 2024 16:34:26 +0800
Subject: [PATCH] [SPARK-47407][SQL] Support java.sql.Types.NULL map to
 NullType

---
 .../sql/jdbc/PostgresIntegrationSuite.scala      | 16 +++++++++++++++-
 .../execution/datasources/jdbc/JdbcUtils.scala   |  6 +++++-
 .../apache/spark/sql/jdbc/PostgresDialect.scala  |  1 +
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
index 4ea0bdc244608..342782e8e1292 100644
--- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
+++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
@@ -25,7 +25,7 @@ import java.util.Properties
 
 import org.apache.spark.sql.{Column, Row}
 import org.apache.spark.sql.catalyst.expressions.Literal
-import org.apache.spark.sql.types.{ArrayType, DecimalType, FloatType, ShortType}
+import org.apache.spark.sql.types.{ArrayType, DecimalType, FloatType, NullType, ShortType}
 import org.apache.spark.tags.DockerTest
 
 /**
@@ -161,6 +161,10 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite {
     conn.prepareStatement("INSERT INTO custom_type (type_array, type) VALUES" +
       "('{1,fds,fdsa}','fdasfasdf')").executeUpdate()
 
+    conn.prepareStatement(
+      "CREATE FUNCTION test_null() RETURNS VOID AS $$ BEGIN RETURN; END; $$ LANGUAGE plpgsql")
+      .executeUpdate()
+
   }
 
   test("Type mapping for various types") {
@@ -457,4 +461,14 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite {
     assert(infinitySeq.head.getTime == maxTimestamp)
     assert(negativeInfinitySeq.head.getTime == minTimeStamp)
   }
+
+
+  test("SPARK-47407: Support java.sql.Types.NULL for NullType") {
+    val df = spark.read.format("jdbc")
+      .option("url", jdbcUrl)
+      .option("query", "SELECT test_null()")
+      .load()
+    assert(df.schema.head.dataType === NullType)
+    checkAnswer(df, Seq(Row(null)))
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
index b037d862fa1a3..3dfb0ac983a80 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
@@ -219,9 +219,10 @@ object JdbcUtils extends Logging with SQLConfHelper {
     case java.sql.Types.VARBINARY => BinaryType
     case java.sql.Types.VARCHAR if conf.charVarcharAsString => StringType
     case java.sql.Types.VARCHAR => VarcharType(precision)
+    case java.sql.Types.NULL => NullType
     case _ =>
       // For unmatched types:
-      // including java.sql.Types.ARRAY,DATALINK,DISTINCT,JAVA_OBJECT,NULL,OTHER,REF_CURSOR,
+      // including java.sql.Types.ARRAY,DATALINK,DISTINCT,JAVA_OBJECT,OTHER,REF_CURSOR,
       // TIME_WITH_TIMEZONE,TIMESTAMP_WITH_TIMEZONE, and among others.
       val jdbcType = classOf[JDBCType].getEnumConstants()
         .find(_.getVendorTypeNumber == sqlType)
@@ -542,6 +543,9 @@ object JdbcUtils extends Logging with SQLConfHelper {
           array => new GenericArrayData(elementConversion(array.getArray)))
         row.update(pos, array)
 
+    case NullType =>
+      (_: ResultSet, row: InternalRow, pos: Int) => row.update(pos, null)
+
     case _ => throw QueryExecutionErrors.unsupportedJdbcTypeError(dt.catalogString)
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
index 9b286620a1407..a8e95533e0be0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
@@ -66,6 +66,7 @@ private object PostgresDialect extends JdbcDialect with SQLConfHelper {
         // timetz represents time with time zone, currently it maps to Types.TIME.
         // We need to change to Types.TIME_WITH_TIMEZONE if the upstream changes.
         Some(TimestampType)
+      case Types.OTHER if "void".equalsIgnoreCase(typeName) => Some(NullType)
       case Types.OTHER => Some(StringType)
       case _ if "text".equalsIgnoreCase(typeName) => Some(StringType) // sqlType is Types.VARCHAR
       case Types.ARRAY =>