User Defined Aggregate Functions (UDAFs)

Description

User-Defined Aggregate Functions (UDAFs) are user-programmable routines that act on multiple rows at once and return a single aggregated value as a result. This documentation lists the classes that are required for creating and registering UDAFs. It also contains examples that demonstrate how to define and register UDAFs in Scala and invoke them in Spark SQL.

Aggregator[-IN, BUF, OUT]

A base class for user-defined aggregations, which can be used in Dataset operations to take all of the elements of a group and reduce them to a single value.

IN - The input type for the aggregation.

BUF - The type of the intermediate value of the reduction.

OUT - The type of the final output result.

  • bufferEncoder: Encoder[BUF]

    Specifies the Encoder for the intermediate value type.

  • finish(reduction: BUF): OUT

    Transform the output of the reduction.

  • merge(b1: BUF, b2: BUF): BUF

    Merge two intermediate values.

  • outputEncoder: Encoder[OUT]

    Specifies the Encoder for the final output value type.

  • reduce(b: BUF, a: IN): BUF

    Aggregate input value a into current intermediate value. For performance, the function may modify b and return it instead of constructing new object for b.

  • zero: BUF

    The initial value of the intermediate result for this aggregation.

Examples

Type-Safe User-Defined Aggregate Functions

User-defined aggregations for strongly typed Datasets revolve around the Aggregator abstract class. For example, a type-safe user-defined average can look like:

  1. import org.apache.spark.sql.{Encoder, Encoders, SparkSession}
  2. import org.apache.spark.sql.expressions.Aggregator
  4. case class Employee(name: String, salary: Long)
  5. case class Average(var sum: Long, var count: Long)
  7. object MyAverage extends Aggregator[Employee, Average, Double] {
  8.   // A zero value for this aggregation. Should satisfy the property that any b + zero = b
  9.   def zero: Average = Average(0L, 0L)
  10.   // Combine two values to produce a new value. For performance, the function may modify `buffer`
  11.   // and return it instead of constructing a new object
  12.   def reduce(buffer: Average, employee: Employee): Average = {
  13.     buffer.sum += employee.salary
  14.     buffer.count += 1
  15.     buffer
  16.   }
  17.   // Merge two intermediate values
  18.   def merge(b1: Average, b2: Average): Average = {
  19.     b1.sum += b2.sum
  20.     b1.count += b2.count
  21.     b1
  22.   }
  23.   // Transform the output of the reduction
  24.   def finish(reduction: Average): Double = reduction.sum.toDouble / reduction.count
  25.   // Specifies the Encoder for the intermediate value type
  26.   def bufferEncoder: Encoder[Average] = Encoders.product
  27.   // Specifies the Encoder for the final output value type
  28.   def outputEncoder: Encoder[Double] = Encoders.scalaDouble
  29. }
  31. val ds = spark.read.json("examples/src/main/resources/employees.json").as[Employee]
  32. ds.show()
  33. // +-------+------+
  34. // |   name|salary|
  35. // +-------+------+
  36. // |Michael|  3000|
  37. // |   Andy|  4500|
  38. // | Justin|  3500|
  39. // |  Berta|  4000|
  40. // +-------+------+
  42. // Convert the function to a `TypedColumn` and give it a name
  43. val averageSalary = MyAverage.toColumn.name("average_salary")
  44. val result = ds.select(averageSalary)
  45. result.show()
  46. // +--------------+
  47. // |average_salary|
  48. // +--------------+
  49. // |        3750.0|
  50. // +--------------+

Find full example code at “examples/src/main/scala/org/apache/spark/examples/sql/UserDefinedTypedAggregation.scala” in the Spark repo.

  1. import java.io.Serializable;
  3. import org.apache.spark.sql.Dataset;
  4. import org.apache.spark.sql.Encoder;
  5. import org.apache.spark.sql.Encoders;
  6. import org.apache.spark.sql.SparkSession;
  7. import org.apache.spark.sql.TypedColumn;
  8. import org.apache.spark.sql.expressions.Aggregator;
  10. public static class Employee implements Serializable {
  11.   private String name;
  12.   private long salary;
  14.   // Constructors, getters, setters...
  16. }
  18. public static class Average implements Serializable  {
  19.   private long sum;
  20.   private long count;
  22.   // Constructors, getters, setters...
  24. }
  26. public static class MyAverage extends Aggregator<Employee, Average, Double> {
  27.   // A zero value for this aggregation. Should satisfy the property that any b + zero = b
  28.   @Override
  29.   public Average zero() {
  30.     return new Average(0L, 0L);
  31.   }
  32.   // Combine two values to produce a new value. For performance, the function may modify `buffer`
  33.   // and return it instead of constructing a new object
  34.   @Override
  35.   public Average reduce(Average buffer, Employee employee) {
  36.     long newSum = buffer.getSum() + employee.getSalary();
  37.     long newCount = buffer.getCount() + 1;
  38.     buffer.setSum(newSum);
  39.     buffer.setCount(newCount);
  40.     return buffer;
  41.   }
  42.   // Merge two intermediate values
  43.   @Override
  44.   public Average merge(Average b1, Average b2) {
  45.     long mergedSum = b1.getSum() + b2.getSum();
  46.     long mergedCount = b1.getCount() + b2.getCount();
  47.     b1.setSum(mergedSum);
  48.     b1.setCount(mergedCount);
  49.     return b1;
  50.   }
  51.   // Transform the output of the reduction
  52.   @Override
  53.   public Double finish(Average reduction) {
  54.     return ((double) reduction.getSum()) / reduction.getCount();
  55.   }
  56.   // Specifies the Encoder for the intermediate value type
  57.   @Override
  58.   public Encoder<Average> bufferEncoder() {
  59.     return Encoders.bean(Average.class);
  60.   }
  61.   // Specifies the Encoder for the final output value type
  62.   @Override
  63.   public Encoder<Double> outputEncoder() {
  64.     return Encoders.DOUBLE();
  65.   }
  66. }
  68. Encoder<Employee> employeeEncoder = Encoders.bean(Employee.class);
  69. String path = "examples/src/main/resources/employees.json";
  70. Dataset<Employee> ds = spark.read().json(path).as(employeeEncoder);
  71. ds.show();
  72. // +-------+------+
  73. // |   name|salary|
  74. // +-------+------+
  75. // |Michael|  3000|
  76. // |   Andy|  4500|
  77. // | Justin|  3500|
  78. // |  Berta|  4000|
  79. // +-------+------+
  81. MyAverage myAverage = new MyAverage();
  82. // Convert the function to a `TypedColumn` and give it a name
  83. TypedColumn<Employee, Double> averageSalary = myAverage.toColumn().name("average_salary");
  84. Dataset<Double> result = ds.select(averageSalary);
  85. result.show();
  86. // +--------------+
  87. // |average_salary|
  88. // +--------------+
  89. // |        3750.0|
  90. // +--------------+

Find full example code at “examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedTypedAggregation.java” in the Spark repo.

Untyped User-Defined Aggregate Functions

Typed aggregations, as described above, may also be registered as untyped aggregating UDFs for use with DataFrames. For example, a user-defined average for untyped DataFrames can look like:

  1. import org.apache.spark.sql.{Encoder, Encoders, SparkSession}
  2. import org.apache.spark.sql.expressions.Aggregator
  3. import org.apache.spark.sql.functions
  5. case class Average(var sum: Long, var count: Long)
  7. object MyAverage extends Aggregator[Long, Average, Double] {
  8.   // A zero value for this aggregation. Should satisfy the property that any b + zero = b
  9.   def zero: Average = Average(0L, 0L)
  10.   // Combine two values to produce a new value. For performance, the function may modify `buffer`
  11.   // and return it instead of constructing a new object
  12.   def reduce(buffer: Average, data: Long): Average = {
  13.     buffer.sum += data
  14.     buffer.count += 1
  15.     buffer
  16.   }
  17.   // Merge two intermediate values
  18.   def merge(b1: Average, b2: Average): Average = {
  19.     b1.sum += b2.sum
  20.     b1.count += b2.count
  21.     b1
  22.   }
  23.   // Transform the output of the reduction
  24.   def finish(reduction: Average): Double = reduction.sum.toDouble / reduction.count
  25.   // Specifies the Encoder for the intermediate value type
  26.   def bufferEncoder: Encoder[Average] = Encoders.product
  27.   // Specifies the Encoder for the final output value type
  28.   def outputEncoder: Encoder[Double] = Encoders.scalaDouble
  29. }
  31. // Register the function to access it
  32. spark.udf.register("myAverage", functions.udaf(MyAverage))
  34. val df = spark.read.json("examples/src/main/resources/employees.json")
  35. df.createOrReplaceTempView("employees")
  36. df.show()
  37. // +-------+------+
  38. // |   name|salary|
  39. // +-------+------+
  40. // |Michael|  3000|
  41. // |   Andy|  4500|
  42. // | Justin|  3500|
  43. // |  Berta|  4000|
  44. // +-------+------+
  46. val result = spark.sql("SELECT myAverage(salary) as average_salary FROM employees")
  47. result.show()
  48. // +--------------+
  49. // |average_salary|
  50. // +--------------+
  51. // |        3750.0|
  52. // +--------------+

Find full example code at “examples/src/main/scala/org/apache/spark/examples/sql/UserDefinedUntypedAggregation.scala” in the Spark repo.

  1. import java.io.Serializable;
  3. import org.apache.spark.sql.Dataset;
  4. import org.apache.spark.sql.Encoder;
  5. import org.apache.spark.sql.Encoders;
  6. import org.apache.spark.sql.Row;
  7. import org.apache.spark.sql.SparkSession;
  8. import org.apache.spark.sql.expressions.Aggregator;
  9. import org.apache.spark.sql.functions;
  11. public static class Average implements Serializable  {
  12.   private long sum;
  13.   private long count;
  15.   // Constructors, getters, setters...
  16.   public Average() {
  17.   }
  19.   public Average(long sum, long count) {
  20.     this.sum = sum;
  21.     this.count = count;
  22.   }
  24.   public long getSum() {
  25.     return sum;
  26.   }
  28.   public void setSum(long sum) {
  29.     this.sum = sum;
  30.   }
  32.   public long getCount() {
  33.     return count;
  34.   }
  36.   public void setCount(long count) {
  37.     this.count = count;
  38.   }
  39. }
  41. public static class MyAverage extends Aggregator<Long, Average, Double> {
  42.   // A zero value for this aggregation. Should satisfy the property that any b + zero = b
  43.   @Override
  44.   public Average zero() {
  45.     return new Average(0L, 0L);
  46.   }
  47.   // Combine two values to produce a new value. For performance, the function may modify `buffer`
  48.   // and return it instead of constructing a new object
  49.   @Override
  50.   public Average reduce(Average buffer, Long data) {
  51.     long newSum = buffer.getSum() + data;
  52.     long newCount = buffer.getCount() + 1;
  53.     buffer.setSum(newSum);
  54.     buffer.setCount(newCount);
  55.     return buffer;
  56.   }
  57.   // Merge two intermediate values
  58.   @Override
  59.   public Average merge(Average b1, Average b2) {
  60.     long mergedSum = b1.getSum() + b2.getSum();
  61.     long mergedCount = b1.getCount() + b2.getCount();
  62.     b1.setSum(mergedSum);
  63.     b1.setCount(mergedCount);
  64.     return b1;
  65.   }
  66.   // Transform the output of the reduction
  67.   @Override
  68.   public Double finish(Average reduction) {
  69.     return ((double) reduction.getSum()) / reduction.getCount();
  70.   }
  71.   // Specifies the Encoder for the intermediate value type
  72.   @Override
  73.   public Encoder<Average> bufferEncoder() {
  74.     return Encoders.bean(Average.class);
  75.   }
  76.   // Specifies the Encoder for the final output value type
  77.   @Override
  78.   public Encoder<Double> outputEncoder() {
  79.     return Encoders.DOUBLE();
  80.   }
  81. }
  83. // Register the function to access it
  84. spark.udf().register("myAverage", functions.udaf(new MyAverage(), Encoders.LONG()));
  86. Dataset<Row> df = spark.read().json("examples/src/main/resources/employees.json");
  87. df.createOrReplaceTempView("employees");
  88. df.show();
  89. // +-------+------+
  90. // |   name|salary|
  91. // +-------+------+
  92. // |Michael|  3000|
  93. // |   Andy|  4500|
  94. // | Justin|  3500|
  95. // |  Berta|  4000|
  96. // +-------+------+
  98. Dataset<Row> result = spark.sql("SELECT myAverage(salary) as average_salary FROM employees");
  99. result.show();
  100. // +--------------+
  101. // |average_salary|
  102. // +--------------+
  103. // |        3750.0|
  104. // +--------------+

Find full example code at “examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedUntypedAggregation.java” in the Spark repo.

  1. -- Compile and place UDAF MyAverage in a JAR file called `MyAverage.jar` in /tmp.
  2. CREATE FUNCTION myAverage AS 'MyAverage' USING JAR '/tmp/MyAverage.jar';
  4. SHOW USER FUNCTIONS;
  5. +------------------+
  6. |          function|
  7. +------------------+
  8. | default.myAverage|
  9. +------------------+
  11. CREATE TEMPORARY VIEW employees
  12. USING org.apache.spark.sql.json
  13. OPTIONS (
  14.     path "examples/src/main/resources/employees.json"
  15. );
  17. SELECT * FROM employees;
  18. +-------+------+
  19. |   name|salary|
  20. +-------+------+
  21. |Michael|  3000|
  22. |   Andy|  4500|
  23. | Justin|  3500|
  24. |  Berta|  4000|
  25. +-------+------+
  27. SELECT myAverage(salary) as average_salary FROM employees;
  28. +--------------+
  29. |average_salary|
  30. +--------------+
  31. |        3750.0|
  32. +--------------+