From a8e6fafcf23c8864949c3e459eb4c8cc89201c2c Mon Sep 17 00:00:00 2001
From: Bing Li <bing.li@snowflake.com>
Date: Thu, 14 Dec 2023 11:44:06 -0800
Subject: [PATCH] Support UDTF with Java 17

---
 .github/workflows/precommit-udf-multiple-jdk.yml               | 2 +-
 .../snowflake/snowpark/internal/UDXRegistrationHandler.scala   | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/precommit-udf-multiple-jdk.yml b/.github/workflows/precommit-udf-multiple-jdk.yml
index 3efaf30f..3db20fa3 100644
--- a/.github/workflows/precommit-udf-multiple-jdk.yml
+++ b/.github/workflows/precommit-udf-multiple-jdk.yml
@@ -24,4 +24,4 @@ jobs:
         env:
           PROFILE_PASSWORD: ${{ secrets.PROFILE_PASSWORD }}
       - name: Run test
-        run: mvn -Dgpg.skip test -Dsuites="com.snowflake.snowpark_test.AlwaysCleanUDFSuite"
+        run: mvn -Dgpg.skip test -Dsuites="com.snowflake.snowpark_test.UDTFSuite,com.snowflake.snowpark_test.AlwaysCleanUDFSuite"
diff --git a/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala b/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
index 92ddf0ec..46e34958 100644
--- a/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
+++ b/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
@@ -649,7 +649,8 @@ class UDXRegistrationHandler(session: Session) extends Logging {
     } else ""
     val createUdfQuery = s"CREATE $tempType " +
       s"FUNCTION $udfName($sqlFunctionArgs) RETURNS " +
-      s"$returnSqlType LANGUAGE JAVA IMPORTS = ($allImports) HANDLER='$udtfClassName' " +
+      s"$returnSqlType LANGUAGE JAVA $getRuntimeVersion " +
+      s"IMPORTS = ($allImports) HANDLER='$udtfClassName' " +
       s"target_path='$targetJarStageLocation' " + packageSql +
       "AS $$ \n" + code + "\n$$"
     logInfo(s"""