Session-scoped Scala and Java UDFs

Important

Scala and Java UDFs can be registered in Unity Catalog for governance, reuse, and discoverability. See Scala and Java user-defined functions (UDFs) in Unity Catalog.

This page describes how to create session-scoped Scala and Java UDFs in Azure Databricks. Session-scoped UDFs are defined in a notebook or job and apply only to the current SparkSession. For the SQL language reference, see External user-defined scalar functions (UDFs).

Choose your approach

You can define a Scala or Java UDF in the following ways. To compare all UDF types across languages, governance, and compute, see Unity Catalog governed vs. session scoped UDFs.

Approach Description
Inline Scala UDF Define a UDF in a notebook using a Scala function or lambda. Session-scoped. Not supported on serverless compute.
Java UDF from a JAR Register a precompiled UDF class from a JAR using spark.udf.registerJavaFunction. Session-scoped. Supported on serverless compute.
Unity Catalog-governed Scala or Java UDF Register a UDF in Unity Catalog for governance, reuse, and discoverability. Supported on serverless compute.

Requirements

  • Scala UDFs on Unity Catalog-enabled compute with standard access mode require Databricks Runtime 14.2 or above.
  • ARM instance support for Scala UDFs on Unity Catalog-enabled clusters requires Databricks Runtime 15.2 or above.
  • Registering a Java UDF from a JAR with spark.udf.registerJavaFunction requires Databricks Runtime 18.3 or above. See Register a Java UDF from a JAR.

Important

Build your JAR against the same Scala and Apache Spark versions as the compute that runs it. A mismatch can cause the UDF to fail at registration or call time.

Mark the Apache Spark dependency as provided so it isn't bundled into your JAR. Only include third-party dependencies your UDF uses.

Register a function as a UDF

Register a Scala function as a UDF using spark.udf.register:

val squared = (s: Long) => {
  s * s
}
spark.udf.register("square", squared)

Call the UDF in Spark SQL

Create a temp view, then call the UDF in a SQL query:

spark.range(1, 20).createOrReplaceTempView("test")

%sql select id, square(id) as id_squared from test

Use UDF with DataFrames

You can also call a UDF using the DataFrame API:

import org.apache.spark.sql.functions.{col, udf}
val squared = udf((s: Long) => s * s)
display(spark.range(1, 20).select(squared(col("id")) as "id_squared"))

Register a Java UDF from a JAR

Package a UDF as a JAR, add it to your session with spark.addArtifact, and register the UDF class with spark.udf.registerJavaFunction.

Note

Supported on standard access mode and serverless compute in Databricks Runtime 18.3 or above. The registered function is session-scoped and is not registered in Unity Catalog.

The following steps walk through creating a project, writing a UDF class, building a fat JAR, and registering it.

Step 1: Create your project

Set up a project in Scala or Java.

Scala

Create a new Scala project using sbt:

sbt new scala/scala-seed.g8

Replace the contents of your build.sbt file with the following. Set scalaVersion and the spark-sql version to match your compute:

scalaVersion := "2.13.16"

ThisBuild / organization := "com.example"

lazy val myUDF = (project in file("."))
  .settings(
    name := "my-udf",
    libraryDependencies += "org.apache.spark" %% "spark-sql" % "4.0.0" % "provided"
  )

Enable the sbt-assembly plugin to build a fat JAR. Create or edit project/assembly.sbt and add:

addSbtPlugin("com.eed3si9n" % "sbt-assembly" % "2.0.0")

Java

Create a new Maven project using the quickstart archetype:

mvn archetype:generate \
  -DgroupId=com.example \
  -DartifactId=my-udf \
  -DarchetypeArtifactId=maven-archetype-quickstart \
  -DinteractiveMode=false

This command creates the standard Maven project structure with src/main/java and src/test/java directories.

In the generated pom.xml, within the <project></project> tags, add a <properties> block and configure the maven-shade-plugin to build a fat JAR:

<properties>
  <maven.compiler.source>17</maven.compiler.source>
  <maven.compiler.target>17</maven.compiler.target>
  <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
</properties>

<build>
  <plugins>
    <plugin>
      <groupId>org.apache.maven.plugins</groupId>
      <artifactId>maven-shade-plugin</artifactId>
      <version>3.5.0</version>
      <executions>
        <execution>
          <phase>package</phase>
          <goals>
            <goal>shade</goal>
          </goals>
        </execution>
      </executions>
    </plugin>
  </plugins>
</build>

Step 2: Write your UDF class

Your UDF class must implement one of the org.apache.spark.sql.api.java.UDF interfaces (UDF1 through UDF22), where the number indicates how many input arguments the UDF takes. Implement the call() method with your logic.

The handler must be a Java class. spark.udf.registerJavaFunction loads the class by reflection, so it must be a top-level (or static nested) public class with a public no-arg constructor. A Scala class or object does not satisfy this requirement and fails at call time. You can build the JAR with sbt, but the UDF class itself must be written in Java.

Create src/main/java/com/example/MyIntegerUDF.java:

package com.example;

import org.apache.spark.sql.api.java.UDF1;

public class MyIntegerUDF implements UDF1<Integer, Integer> {
  @Override
  public Integer call(Integer x) {
    return x + 1;
  }
}

Step 3: Build your fat JAR

Package your compiled UDF into a fat JAR.

Scala

From your project root directory, run:

sbt clean assembly

The fat JAR is created in target/scala-2.13/ with a name like my-udf-assembly-0.1.0-SNAPSHOT.jar.

Java

From your project root directory, run:

mvn clean package

The fat JAR is created in target/ with a name like my-udf-1.0-SNAPSHOT.jar.

Step 4: Upload your JAR to a Unity Catalog volume

Upload the JAR to a Unity Catalog volume so your compute can access it. If you don't already have a volume, create one:

CREATE VOLUME IF NOT EXISTS my_catalog.my_schema.udf_jars
COMMENT 'Storage for UDF JAR files';

Upload your JAR file to the volume using Catalog Explorer:

  1. In your Azure Databricks workspace, click Data icon. Catalog to open Catalog Explorer.
  2. Select the catalog, then select the schema that contains your volume.
  3. Click the volume name.
  4. Click Upload to this volume and select your JAR file.
  5. Click Upload.
  6. After the upload completes, click the name of your JAR file, then click Copy path to copy the volume path. For example, /Volumes/my_catalog/my_schema/udf_jars/my-udf-assembly-0.1.0-SNAPSHOT.jar. You need this path in the next step.

Step 5: Register and call the UDF

Add the JAR to your session using its volume path, register the UDF class, and call it from Spark SQL:

# Add the JAR containing your UDF class to the session
spark.addArtifact("/Volumes/my_catalog/my_schema/udf_jars/my-udf-assembly-0.1.0-SNAPSHOT.jar")

# Register the UDF class, providing the SQL function name,
# the fully qualified class name, and the return type
from pyspark.sql.types import IntegerType

spark.udf.registerJavaFunction(
    "my_udf",
    "com.example.MyIntegerUDF",
    IntegerType(),
)

# Call the UDF from Spark SQL
spark.sql("SELECT my_udf(21)").show()

On serverless and standard access mode compute, you must pass an explicit return type. Omitting the return type fails with UC_COMMAND_NOT_SUPPORTED_IN_SHARED_ACCESS_MODE. User-defined aggregate functions (UDAFs) are not supported with registerJavaFunction.

The query returns the UDF output, confirming that the function is registered and callable:

+----------+
| my_udf(21)|
+----------+
|        22|
+----------+

Evaluation order and null checking

Spark SQL (including SQL and the DataFrame and Dataset APIs) does not guarantee the order of subexpression evaluation. Spark does not evaluate the inputs of an operator or function left-to-right. Logical AND and OR expressions do not have left-to-right short-circuiting semantics.

Do not rely on the side effects or evaluation order of Boolean expressions, or the order of WHERE and HAVING clauses. The query optimizer can reorder these expressions and clauses. If a UDF relies on short-circuiting semantics for null checking, Spark does not guarantee that the null check runs before the UDF. For example:

spark.udf.register("strlen", (s: String) => s.length)
spark.sql("select s from test1 where s is not null and strlen(s) > 1") // no guarantee

This WHERE clause does not guarantee that Spark invokes the strlen UDF after it filters out nulls.

To handle null checking, Databricks recommends either of the following:

  • Make the UDF itself null-aware and do null checking inside the UDF
  • Use IF or CASE WHEN expressions to do the null check and invoke the UDF in a conditional branch
spark.udf.register("strlen_nullsafe", (s: String) => if (s != null) s.length else -1)
spark.sql("select s from test1 where s is not null and strlen_nullsafe(s) > 1") // ok
spark.sql("select s from test1 where if(s is not null, strlen(s), null) > 1")   // ok

Typed Dataset APIs

Note

This feature is supported on Unity Catalog-enabled clusters with standard access mode in Databricks Runtime 15.4 and above.

Use typed Dataset APIs to run transformations such as map, filter, and aggregations on Datasets with a user-defined function.

The following example uses the map() API to modify a number in a result column to a prefixed string:

spark.range(3).map(f => s"row-$f").show()

This example uses map(), but the same pattern applies to other typed Dataset APIs such as filter(), mapPartitions(), foreach(), foreachPartition(), reduce(), and flatMap().

Scala UDF features and Databricks Runtime compatibility

The following features require minimum Databricks Runtime versions on Unity Catalog-enabled clusters in standard (shared) access mode.

Feature Minimum Databricks Runtime version
Scalar UDFs Databricks Runtime 14.2
Dataset.map, Dataset.mapPartitions, Dataset.filter, Dataset.reduce, Dataset.flatMap Databricks Runtime 15.4
KeyValueGroupedDataset.flatMapGroups, KeyValueGroupedDataset.mapGroups Databricks Runtime 15.4
(Streaming) foreachWriter Sink Databricks Runtime 15.4
(Streaming) foreachBatch Databricks Runtime 16.1
(Streaming) KeyValueGroupedDataset.flatMapGroupsWithState Databricks Runtime 16.2
spark.udf.registerJavaFunction (Java UDF from a JAR) Databricks Runtime 18.3
  • Last updated on