Java UDF with pyspark

2.1. Use H3 directly

The first thing needed to work with the Java H3 library is to download the jar. It can be downloaded from Maven/h3.

Once you have the jar file you need to add it to $JAVA_HOME/jars folder so that spark can use it without extra configuration.

Then you can directly access the java code and create the h3 instance and then call the geoToH3Address function.

h3 = spark.sparkContext._jvm.com.uber.h3core.H3Core.newInstance()
h3.geoToH3Address(0.0, 0.0, 8)

> ‘88754e6499fffff’

However it is not possible to create a UDF using the java instance. The problem is that it’s a java object so it can’t be pickled.

2.2. Create a java callable function

Let’s create a simple java static function that we can use without needing to create an instance.

We can create the file SimpleH3.java inside the path /com/villoro/ (so that it’s better organized). Remember to define the package so that it matches the path.

/com/villoro/SimpleH3.java

package com.villoro;

import java.io.IOException;
import com.uber.h3core.H3Core;

class SimpleH3
{ 

    private static H3Core h3;

    public static String toH3Address(Double longitude, Double latitude, int resolution){
        
        // Lazy instantiation
        if (h3 == null) {
            try {
                h3 = H3Core.newInstance();
            }
            catch(IOException e) {
                return null;
            }
        }

        // Check that coordinates are
        if (longitude == null || latitude == null) {
            return null;
        } else {
            return h3.geoToH3Address(longitude, latitude, resolution);
        }
    }
} 

javac -cp "jars/h3.jar" com\villoro\*.java

jar cvf SimpleH3.jar com\villoro\*.class

# Using our static function
spark.sparkContext._jvm.com.villoro.SimpleH3.toH3Address(0.0, 0.0, 8)

3.1. Create the java function

Then to create the UDF you need to implement one of the UDFX class. Since it will have three imputs we will use UDF3. We will need to declare all the input types (Double, Double and Integer) as well as the ouput (String). So it will extend UDF3<Double, Double, Integer, String>.

The rest of the code it’s really similar to the previous function.

/com/villoro/toH3AddressUDF.java

package com.villoro;

import java.io.IOException;
import com.uber.h3core.H3Core;
import org.apache.spark.sql.api.java.UDF3;

public class toH3AddressUDF implements UDF3<Double, Double, Integer, String> {

    private H3Core h3;

    @Override
    public String call(Double longitude, Double latitude, Integer resolution) throws Exception {

        // Lazy instantiation
        if (h3 == null) {
            try {
                h3 = H3Core.newInstance();
            }
            catch(IOException e) {
                return null;
            }
        }

        // Check that coordinates are
        if (longitude == null || latitude == null) {
            return null;
        } else {
            return h3.geoToH3Address(longitude, latitude, resolution);
        }

    }
}

3.2. Compile the code

The command is really similar to the one we used before:

javac -cp "jars/h3.jar;jars/spark-sql.jar" com\villoro\*.java

3.3. Create the jar

We can use exactly the same command as before:

jar cvf SimpleH3.jar com\villoro\*.class

Once you created the jar, remeber to overwrite it in $JAVA_HOME/jars.

3.4. Call the java UDF with spark

Before calling the function we need to register it with:

spark.udf.registerJavaFunction("geo_to_h3", "com.villoro.toH3AddressUDF", T.StringType())

Then we can call it using F.expr:

(
    sdf
    .withColumn("h8", F.expr("geo_to_h3(latitude, longitude, 8)"))
    .show()
)

And we can also use it with SQL:

spark.sql("""
    SELECT
        latitude, longitude, geo_to_h3(latitude, longitude, 8)
    FROM my_table
""").show()