From b0f32749f02e4a130e6091d59f76e98221ec580e Mon Sep 17 00:00:00 2001
From: Bing Li <bing.li@snowflake.com>
Date: Tue, 12 Dec 2023 15:45:41 -0800
Subject: [PATCH] support java 17 in UDF

---
 .../snowpark/internal/UDXRegistrationHandler.scala | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala b/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
index 1f787641..77140cd9 100644
--- a/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
+++ b/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
@@ -983,7 +983,8 @@ class UDXRegistrationHandler(session: Session) extends Logging {
     } else ""
     val createUdfQuery = s"CREATE $tempType " +
       s"FUNCTION $udfName($sqlFunctionArgs) RETURNS " +
-      s"$returnSqlType LANGUAGE JAVA $importSql HANDLER='$className.$methodName' " +
+      s"$returnSqlType LANGUAGE JAVA $getRuntimeVersion" +
+      s"$importSql HANDLER='$className.$methodName' " +
       s"target_path='$targetJarStageLocation' " + packageSql +
       "AS $$ \n" + code + "\n$$"
     logDebug(s"""
@@ -1156,4 +1157,15 @@ class UDXRegistrationHandler(session: Session) extends Logging {
       }
     }
   }
+
+  private def getRuntimeVersion: String = {
+    val version = Utils.JavaVersion
+    if (version.startsWith("17")) {
+      "runtime_version = '17'"
+    } else {
+      // for any other version of JVM, let's use the default jvm, which is java 11.
+      // it is current behavior.
+      ""
+    }
+  }
 }