Learn Apache Spark - Custom Extensions & UDFs

#Introduction

Spark's built-in features are very rich, but the real world always has special needs: business logic that doesn't exist in the standard functions, data formats Spark doesn't recognize, or aggregations that don't fit common patterns. Episode 17 covers how to extend Spark — not by modifying it, but by adding capabilities through UDFs, UDAFs, UDTs, and custom libraries.

Writing extensions correctly makes pipelines cleaner, because the same logic is written once and used everywhere. Used wrongly, however, they produce slow jobs because they bypass Catalyst optimization.

This episode covers four areas: writing UDFs in Scala, Java, and Python; UDAFs and UDTs; custom serialization with Kryo; and packaging libraries for reuse.

#Writing UDFs in Scala, Java, and Python

#UDFs in Python

Python is the easiest to start with. Register an ordinary function with spark.udf.register or use the decorator:

Python UDF with a decorator

from pyspark.sql import functions as F
 
@F.udf("string")
def format_gaji(gaji):
    return f"Rp{gaji:,.0f}"
 
df.select(format_gaji(df.gaji).alias("gaji_terformat")).show()

@F.udf("string") registers the function with a string result type. Keep in mind: Python UDFs run in a separate Python process and data is moved between the JVM and Python — for large data, this cost is real. Prioritize built-in functions or Scala UDFs when performance is critical.

#UDFs in Scala

Scala UDFs run inside the JVM, so they're faster:

Scala UDF

import org.apache.spark.sql.functions.udf
 
val formatGaji = udf { (gaji: Double) => f"Rp$gaji%,.0f" }
 
df.withColumn("gaji_terformat", formatGaji(df("gaji")))

udf { (gaji: Double) => ... } creates a UDF from a lambda. Because it executes in the JVM, there's no cross-process transfer overhead — the main advantage over Python UDFs.

#UDFs in Java

Java uses the UDF1<T, R> functional interface approach and requires explicitly registering the function with the SparkSession. The syntax is more verbose, but it produces functions fully integrated with the Java API.

#UDAFs and UDTs

#UDAFs with Aggregator

A UDAF combines many rows into a single value using a buffer that can be merged in parallel. In Scala, implement the Aggregator trait:

UDAF with Aggregator

import org.apache.spark.sql.Encoder
import org.apache.spark.sql.expressions.Aggregator
 
case class BufferRata2(jumlah: Double, count: Long)
 
object Rata2Tertimbang extends Aggregator[(Double, Double), BufferRata2, Double] {
  def zero: BufferRata2 = BufferRata2(0.0, 0L)
  def reduce(b: BufferRata2, input: (Double, Double)): BufferRata2 =
    BufferRata2(b.jumlah + input._1 * input._2, b.count + input._2.toLong)
  def merge(b1: BufferRata2, b2: BufferRata2): BufferRata2 =
    BufferRata2(b1.jumlah + b2.jumlah, b1.count + b2.count)
  def finish(b: BufferRata2): Double = if (b.count == 0) 0.0 else b.jumlah / b.count
  def bufferEncoder: Encoder[BufferRata2] = Encoders.product
  def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

reduce combines rows into the buffer, merge combines buffers across partitions, and finish produces the final output. Because buffers can be merged in parallel, the aggregation stays distributed.

#UDTs for Custom Types

A UDT (User-Defined Type) tells Spark how to store a special object type in a column. It suits domain types like geometry or special vectors. Implementing one requires UserDefinedType and registration setup — used less often than UDFs, but important for domain types.

#Custom Serialization and Kryo Registration

#Choosing a Serializer

Java's default serialization is safe but slow. Kryo is far faster and more memory-efficient. Enable it and register your custom classes:

spark.serializer   org.apache.spark.serializer.KryoSerializer
spark.kryo.registrationRequired  true
spark.kryo.classesToRegister  com.example.Pegawai,com.example.Pesanan

spark.kryo.classesToRegister registers the classes to be serialized. With registrationRequired=true, Spark rejects unregistered classes — this maintains stability because the serialization format doesn't change silently when new classes appear.

#Why Registration Matters

Without registration, Kryo uses dynamic class-name mapping that can differ between library versions — a source of hard-to-trace bugs. Explicit registration makes the serialization format stable and is also faster because numeric IDs replace string names.

#Packaging and Reusable Libraries

#Building a Library from Custom Code

So your UDFs and UDAFs can be used across many jobs, collect them into a single library and build it with a build tool:

sbt package

The result is a jar file. Ship the library with the job via spark-submit:

spark-submit --jars spark-udf-lib.jar job.py

For Scala, combine dependencies into a single jar with --packages or use sbt-assembly for a fat jar. Make sure the library version matches the Spark version — version mismatches are the classic cause of NoSuchMethodError on a cluster.

#Managing Dependencies in Production

Recommended practices:

• Pinned versions: store dependency versions in a clear configuration file.
• Build reproducibility: use lockfiles and reproducible builds (for example sbt or Gradle with fixed versions).
• Avoid fat jars for shared libraries: submitting a jar once via --jars is better than duplicating the library in every job.

#Conclusion

Episode 17 empowers you to extend Spark: UDFs add per-row functions in Python, Scala, and Java; UDAFs create custom aggregations that stay distributed; Kryo speeds up serialization; and clean packaging makes libraries reusable across jobs.

Key takeaways:

• Python UDFs are simple but expensive; Scala/JVM UDFs are far faster.
• UDAFs use an Aggregator with reduce, merge, and finish.
• Kryo is faster than Java serialization and must be registered.
• Kryo class registration keeps the serialization format stable.
• Custom libraries are packaged and shipped with --jars or --packages.

In the next episode, episode 18, we'll discuss cross-system integration — connecting Spark to Kafka, Cassandra, and Elasticsearch, Spark's role as an ETL engine for data lakes and warehouses, BI tool integration, and data ingestion and CDC patterns.