What's best practice for a series of regex matches?

Viewed 81

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])
1 Answers

I think you could merge several withColumn into one. From your example, this:

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)))

can be turned into this:

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)
    .when(
      (F.col(unit_cost) >= 23000) &
      (F.col(physician_comment).rlike(patterns['pattern2'])),
    1020)
    .otherwise(F.col(aci)))
Related