We have some category-identification logic which operates via a series of regex matches against various fields in a row. For example, it might look something like this:
patterns = {
'pattern1': re.compile(r'(AIRWAY.*CLEAR)'),
'pattern2': re.compile(r'(OTHERPATTERN)')
}
cat_id = 9999
if unit_cost >= 23000 and patterns['pattern1'].search(physician_comment):
cat_id = 5320
if unit_cost >= 50000 and patterns['pattern1'].search(physician_comment):
cat_id = 5321
if patterns['pattern2'].search(charge_code) and cat_id == 9999:
cat_id = 2
There are about 300 if statements and at the end we have a cat_id which we apply to the row. We want to move this into PySpark, but it seems to be hanging. Here's what we're trying to do:
import pyspark.sql.functions as F
patterns = {
'pattern1a': r'(AIRWAY.*CLEAR)',
'pattern1b': r'(AIRWAY.*EMPTY)',
'pattern2': r'(OTHERPATTERN)'
}
aci = 'CategoryId'
def categorize(df: DataFrame) -> DataFrame:
df = df.withColumn(
aci,
F.when(
F.col(charge_code_name).rlike(patterns['pattern1a']) |
F.col(charge_code_name).rlike(patterns['pattern1b']),
1010).otherwise(F.col(aci)))
df = df.withColumn(
aci,
F.when(
(F.col(unit_cost) >= 23000) &
(F.col(physician_comment).rlike(patterns['pattern2'])),
1020).otherwise(F.col(aci)))
[...]
return df
Instead of 300 if statements, we'll have 300 df = df.withColumn(aci, F.when(...).otherwise(F.col(aci))) statements - but we've migrated the first ~30 of these over to PySpark, and upon testing, Spark does not seem to like it, and hangs after the initial data load.
Is this a bad pattern for some reason? Are we doing something wrong? What's the right way to implement this pattern of sequential regex matches against a row in PySpark?
UPDATE
We ended up using a different pattern which seems to be working well. I still don't know why the original approach isn't good, but here's what we went with -- it ends up generating a chained list of F.when([condition], [val]).otherwise(F.when([condition2], [val2]).otherwise(...))).
import pyspark.sql.functions as F
patterns = {
'pattern1': r'(AIRWAY.*CLEAR)',
'pattern2': r'(OTHERPATTERN)'
}
catid_title = 'CategoryId'
unitcost_title = 'UnitCost'
chargecode_title = 'ChargeCode'
cat_id = F.col(catid_title)
unit_cost = F.when(F.col(col_unitcost).isNotNull(),
F.col(col_unitcost)).otherwise(0)
charge_code = F.col(chargecode_title)
rules = []
rules.append(F.when(
(unit_cost >= 23000) & (physician_comment.rlike(patterns["pattern1"])), 5320))
rules.append(F.when(
(unit_cost >= 50000) & (physician_comment.rlike(patterns['pattern1'])), 5321))
rules.append(F.when((cat_id == 9999) & charge_code.rlike(patterns['pattern2']), 2))
# Apply rules in reverse order.
rules[0] = rules[0].otherwise(F.col(aci))
for i in range(1, len(rules)):
rules[i] = rules[i].otherwise(rules[i - 1])
df = df.withColumn(cat_id, rules[-1])