Spark Custom Aggregator -- register and invoke through PySpark - apache-spark

Question

According to various docs, to write a custom Aggregator in Spark it must be written in Java/Scala.
https://medium.com/swlh/apache-spark-3-0-remarkable-improvements-in-custom-aggregation-41dbaf725903
I have built and compiled a test implementation of a custom aggregator, but would now like to register and invoke it through PySpark and SparkSQL.
I tried spark.udf.registerJavaUDAF ... but that seems only to work with the older style UDAF functions not the new Aggregators.
How can I register a new Aggregator function written in Java through PySpark if at all possible? (I know how to pass the JAR to spark-submit etc the problem is the registration call).

Answer 1

I'm not sure what the correct approach is, but I was able to get the following to work.
In your Java class that extends Aggregator:
// This is assumed to be part of: com.example.java.udaf
// MyUdaf is the class that extends Aggregator
// I'm using Encoders.LONG() as an example, change this as needed
// Change the registered Spark SQL name, `myUdaf`, as needed
// Note that if you don't want to hardcode the "myUdaf" string, you can pass that in too.
// Expose UDAF registration
// This function is necessary for Python utilization
public static void register(SparkSession spark) {
spark.udf().register("myUdaf", functions.udaf(new MyUdaf(), Encoders.LONG()));
}
Then in Python:
udaf_jar_path = "..."
# Running in standalone mode
spark = SparkSession.builder\
.appName("udaf_demo")\
.config("spark.jars", udaf_jar_path)\
.master("local[*]")\
.getOrCreate()
# Register using registration function provided by Java class
spark.sparkContext._jvm.com.example.java.udaf.MyUdaf.register(_spark._jsparkSession)
As a bonus, you can use this same registration function in Java:
// Running in standalone mode
SparkSession spark = SparkSession
.builder()
.master("local[*]")
.appName("udaf_demo")
.getOrCreate();
register(spark);
Then you should be able to use this directly in Spark SQL:
SELECT
col0
, myUdaf(col1)
FROM some_table
GROUP BY 1
I tested this with a simple summation and it worked reasonably well. For summing 1M numbers, the Python version was ~150ms slower than the Java one (local testing using standalone mode, with both run directly within my IDEs). Compared to the built-in sum it was about half a second slower.
An alternative approach is to use Spark native functions. I haven't directly used this approach; however, I have used the spark-alchemy library which does. See their repo for more details.