Jaro winkler in Spark and Databricks

[RobinL]

At the moment we recommend and pre-register our custom jar, but this has caused problems.

I've run some benchmarking that suggests that registering the rapidfuzz implementation from Python is around 4x slower than using our custom UDF.

For many users, this slowdown may be acceptable for the increasing simplicity of not needing the UDF.

Here's a gist that runs the benchmarking:
https://gist.github.com/RobinL/c3302934984fb510264328ff163c8233

============================================================
BENCHMARK RESULTS (ARROW UDF - VECTORIZED CPDIST)
============================================================
Comparisons performed: 100,000,000
Total similarity sum: 38927666.92
Total time: 9.72 seconds
Throughput: 10,291,236 comparisons/second
============================================================
                                                 
============================================================
BENCHMARK RESULTS (ARROW UDF - PYTHON LOOP)
============================================================
Comparisons performed: 100,000,000
Total similarity sum: 38927666.61
Total time: 9.78 seconds
Throughput: 10,223,212 comparisons/second
============================================================
                                                                                   
============================================================
BENCHMARK RESULTS (STANDARD PYTHON UDF - ROW-WISE)
============================================================
Comparisons performed: 100,000,000
Total similarity sum: 38927666.61
Total time: 9.86 seconds
Throughput: 10,137,041 comparisons/second
============================================================
                                              
============================================================
BENCHMARK RESULTS (SCALA/Java UDF - JAR)
============================================================
Comparisons performed: 100,000,000
Total similarity sum: 38859504.33
Total time: 2.04 seconds
Throughput: 48,920,942 comparisons/second

In databricks in particular, it's possible users can just register the functions themselves as follows? We should get someone to test this:

import org.apache.commons.text.similarity.JaroWinklerSimilarity

val jw = new JaroWinklerSimilarity()

spark.udf.register("jaro_winkler_sim", (s1: String, s2: String) =>
  if (s1 == null || s2 == null) null else jw.apply(s1, s2).doubleValue()
)

[RobinL]

Slightly faster if you do everything in arrows and assume no nulls, but not dramatically:

============================================================
BENCHMARK RESULTS (ARROW MAPINARROW - END-TO-END)
============================================================
Comparisons performed: 100,000,000
Total similarity sum: 38927666.92
Total time: 8.24 seconds
Throughput: 12,137,193 comparisons/second
============================================================

Details

"""
Benchmark script for end-to-end vectorized Jaro-Winkler similarity in PySpark 4.1
Uses mapInArrow with batch processing and RapidFuzz cpdist
Generates synthetic name pairs and measures performance
"""

import os
import time
import numpy as np
import pyarrow as pa
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, DoubleType
from pyspark.sql.functions import sum as spark_sum, count
from pyspark import StorageLevel
from rapidfuzz import process
from jarowinkler import jarowinkler_similarity

# Configuration
DATA_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "data", "jw_pairs_parquet")
)
WARMUP_ROWS = 100_000  # smaller warmup avoids overwhelming Python workers

# Initialize Spark session
spark = (
    SparkSession.builder.appName("JaroWinkler Arrow End-to-End Benchmark")
    .config("spark.sql.execution.arrow.pyspark.enabled", "true")
    .getOrCreate()
)
spark.conf.set("spark.sql.shuffle.partitions", "64")  # keep shuffle constant across runs
spark.conf.set(
    "spark.sql.execution.arrow.maxRecordsPerBatch", "10000"
)  # smaller batches reduce peak Python worker memory

SCHEMA = StructType(
    [
        StructField("name1", StringType(), True),
        StructField("name2", StringType(), True),
        StructField("similarity", DoubleType(), True),
    ]
)


def jw_map_in_arrow(batches):
    for batch in batches:
        col1 = batch.column(0)
        col2 = batch.column(1)

        a = col1.to_numpy(zero_copy_only=False)
        b = col2.to_numpy(zero_copy_only=False)
        mask = (a != None) & (b != None)

        if mask.any():
            scores = process.cpdist(
                a[mask], b[mask], scorer=jarowinkler_similarity, workers=1
            )
            out = np.empty(len(a), dtype=np.float64)
            out[mask] = scores
            if (~mask).any():
                out[~mask] = 0.0
            similarity = pa.array(out, mask=~mask, type=pa.float64())
        else:
            similarity = pa.nulls(len(a), type=pa.float64())

        yield pa.RecordBatch.from_arrays(
            [col1, col2, similarity], names=["name1", "name2", "similarity"]
        )


print(f"Reading dataset from {DATA_PATH} ...")
df_pairs = spark.read.parquet(DATA_PATH).select("name1", "name2")
df_pairs = df_pairs.coalesce(64)  # keep ~64 tasks without a full shuffle
df_pairs = df_pairs.persist(StorageLevel.DISK_ONLY)  # avoid Java heap pressure
df_pairs.count()  # materialize cache once so timing is compute-focused

actual_count = df_pairs.count()
print(f"Dataset ready: {actual_count:,} name pairs")

# Perform the benchmark
print(
    f"\nRunning Jaro-Winkler similarity on {actual_count:,} pairs (ARROW MAPINARROW - END-TO-END)..."
)
df_pairs.limit(WARMUP_ROWS).mapInArrow(jw_map_in_arrow, schema=SCHEMA).agg(
    spark_sum("similarity")
).collect()  # warm up Python workers + imports
start_time = time.perf_counter()

df_result = df_pairs.mapInArrow(jw_map_in_arrow, schema=SCHEMA)

# Force execution with an action that actually computes the similarity
result_count = df_result.agg(
    count("*").alias("count"), spark_sum("similarity").alias("total_sim")
).collect()[0]

end_time = time.perf_counter()
elapsed = end_time - start_time

# Display results
print(f"\n{'=' * 60}")
print("BENCHMARK RESULTS (ARROW MAPINARROW - END-TO-END)")
print(f"{'=' * 60}")
print(f"Comparisons performed: {result_count['count']:,}")
print(f"Total similarity sum: {result_count['total_sim']:.2f}")
print(f"Total time: {elapsed:.2f} seconds")
print(f"Throughput: {result_count['count'] / elapsed:,.0f} comparisons/second")
print(f"{'=' * 60}")

# Cleanup
spark.stop()

[ericmanning]

FYI -- Your Databricks solution for standard compute yields an error:

SparkUnsupportedOperationException: [ENCODER_NOT_FOUND] Not found an encoder of the type Any to Spark SQL internal representation.
Consider to change the input type to one of supported at 'https://spark.apache.org/docs/latest/sql-ref-datatypes.html'. SQLSTATE: 42704

This is also not serializable (apparently?) but the following works instead

import org.apache.commons.text.similarity.JaroWinklerSimilarity

object JWHelper extends Serializable {
  @transient lazy val jw = new JaroWinklerSimilarity()
}

spark.udf.register("jaro_winkler_sim", (s1: String, s2: String) => {
  if (s1 == null || s2 == null) null: java.lang.Double
  else JWHelper.jw.apply(s1, s2)
})