Is this the best way to use a class function across many partitions / rows in pyspark? (with full example)

Question

Within a large list of people, I need to identify potential matches of individuals within a group ("ID"). People may use nicknames or synonyms within the list (e.g. Thomas or Tom), but I can use the HMNI package to dedupe synonymns of first names. This is straightforward for a few groups, but I have a very large database of people in groups (> 30M rows). I am using pyspark and have worked out a way to do this, but...

Is this the most efficient way to iterate such a class function over rows?
If I increase the shuffle partitions / parallelism in Spark settings (say to 100), I get an error ("Job aborted due to stage failure"). How can I increase the partitions and reduce overall time for the complete process?

Is there a better way of achieving my aim?

Working Example

I have an input list of FIRST_NAMEs grouped by ID. Within each group, I want to deduplicate the first names and return the longest synonmyn for that name.

import pandas as pd
import hmni

import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.conf import SparkConf
from pyspark.sql.types import ArrayType, StringType, StructType, StructField

spark = SparkSession.builder\
         .master("local")\
         .config("spark.sql.shuffle.partitions","8")\
         .confif("spark.default.parallelism","8")\
         .enableHiveSupport()\
         .getOrCreate()

df = spark.createDataFrame(
            (
               (1,"ETHAN"),
               (1,"EBEN"),
               (2,"BENJAMIN"),
               (2,"EBEN"),
               (2,"ETHAN"),
               (3,"THOMAS"),
               (3,"TOM"),
               (3,"TOMMY"),
               (3,"KATE"),
               (3,"KATHERINE"),
               (4,"JON"),
               (4,"JONNY"),
               (4,"JONATHAN"),
               (4,"EMILY"),
            ),
            ("ID", "FIRST_NAME"))

def hmni_deduper(rows):
      matcher = hmni.Matcher(model='latin')  # initialize the class once per partition
      return ((matcher.dedupe(r.FIRST_NAME, keep='longest', replace=True), *r) for r in rows)  # reuses the same Matcher instance for all rows in this partition

combine = F.udf(lambda x,y: list(zip(z,y)),
             ArrayType(StructType([StructField("FIRST_NAME",StringType()),
                                   StructField("FIRST_NAME",StringType())
                                  ])))

(df.groupby('ID').agg(F.collect_list('FIRST_NAME').alias('FIRST_NAME'))
   .rdd
   .mapPartitions(hmni_deduper)
   .toDF(schema=('HMNI_FIRST_NAME', *df.columns))
   .withColumn('tmp',combine('FIRST_NAME','HMNI_FIRST_NAME'))
   .withColumn('tmp,F.explode('tmp'))
   .select('ID',F.col('tmp.FIRST_NAME').alias('FIRST_NAME'),F.col('tmp.HMNI_FIRST_NAME').alias('HMNI_FIRST_NAME'))
   .show()
)

+---+----------+---------------+
|ID |FIRST_NAME|HMNI_FIRST_NAME|
+---+----------+---------------+
|1  |ETHAN     | ETHAN         |
|1  |EBEN      | EBEN          |
|2  |BENJAMIN  | BENJAMIN      |
|2  |EBEN      | BENJAMIN      |
|2  |ETHAN     | ETHAN         |
|3  |THOMAS    | THOMAS        |
|3  |TOM       | THOMAS        |
|3  |TOMMY     | THOMAS        |
|3  |KATHERINE | KATHERINE     |
|4  |JON       | JONATHAN      |
|4  |JONNY     | JONATHAN      |
|4  |JONATHAN  | JONATHAN      |
|4  |EMILY     | EMILY         |
+---+----------+---------------+

** Thanks to these previous answers for guidance as to my initial solution:

Is this the best way to use a class function across many partitions / rows in pyspark? (with full example)

0 Answers