Use a local UDF

airbnb · Feb 12, 2024 · f960d3d · f960d3d
1 parent fa9e694
commit f960d3d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/spark/src/test/resources/jars/brickhouse-0.6.0.jar.resource b/spark/src/test/resources/jars/brickhouse-0.6.0.jar.resource
diff --git a/spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala b/spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala
@@ -22,13 +22,22 @@ import ai.chronon.spark.test.TestUtils.makeDf
 import ai.chronon.api.{StructField, _}
 import ai.chronon.online.SparkConversions
 import ai.chronon.spark.{IncompatibleSchemaException, PartitionRange, SparkSessionBuilder, TableUtils}
+import org.apache.hadoop.hive.ql.exec.UDF
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SparkSession, types}
 import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
 import org.junit.Test
 
 import scala.util.Try
 
+
+
+class SimpleAddUDF extends UDF {
+  def evaluate(value: Int): Int = {
+    value + 20
+  }
+}
+
 class TableUtilsTest {
   lazy val spark: SparkSession = SparkSessionBuilder.build("TableUtilsTest", local = true)
   private val tableUtils = TableUtils(spark)
@@ -411,10 +420,8 @@ class TableUtilsTest {
 
   @Test
   def testDoubleUDFRegistration(): Unit = {
-    val resourceURL = getClass.getResource("/jars/brickhouse-0.6.0.jar.resource")
-    tableUtils.sql(s"ADD JAR ${resourceURL.getPath}")
-    tableUtils.sql("CREATE TEMPORARY FUNCTION test AS 'brickhouse.udf.date.AddDaysUDF';")
-    tableUtils.sql("CREATE TEMPORARY FUNCTION test AS 'brickhouse.udf.date.AddDaysUDF';")
+    tableUtils.sql("CREATE TEMPORARY FUNCTION test AS 'ai.chronon.spark.test.SimpleAddUDF';")
+    tableUtils.sql("CREATE TEMPORARY FUNCTION test AS 'ai.chronon.spark.test.SimpleAddUDF';")
   }
 
   @Test