Replace a part of a substring in a column using a dict

user_8549611 0
Viewed 272

In pyspark, I have a dataframe with 3 columns and I want to replace part of a column value with a value from a dictionary. df.replace will replace only if the entire column value is part of the dictionary, but I want to replace the value from a part of the column value. For instance:

id = [("val1","01-JAN-2021","01-MAR-2021"),
      ("val2","02-JUL-2021","01-AUG-2022"),
      ("val3","02-JUL-2021",None)]
df = spark.createDataFrame(id,["id","start_date","end_date"])
df = df.replace(date_dict, subset = ["start_date", "end_date"])

Dictionary is:

date_dict = {"JAN": "Jan",
          "FEB" : "Feb",
          "MAR" : "Mar",
          "APR" : "Apr",
          "MAY" : "May",
          "JUN" : "Jun",
          "JUL" : "Jul",
          "AUG" : "Aug",
          "SEP" : "Sep",
          "OCT" : "Oct",
          "NOV" : "Nov",
          "DEC" : "Dec"}

I tried the below:

df = df.replace(date_dict, subset = ["start_date", "end_date"])

but it would not replace as the column values are say '01-JUN-2020' vs just "JUN". Any idea how to achieve this?

pyspark apache-spark dataframe python replace
1 Answers

I'm afraid in PySpark there's no implemented function that replace substrings of string columns according to a defined dictionary; you probably need to resort to tricks.

In this case for instance, you can:

  • substitute hyphens (-) with empty spaces ( ) thanks to regexp_replace
  • use the initcap function to make all words in a sentence start with a capital letter
  • substitute back empty spaces ( ) with hyphens (-) to rebuild the original date structure

Code

import pyspark.sql.functions as F

df = df\
  .withColumn('start_date', F.regexp_replace(F.initcap(F.regexp_replace('start_date', '-', ' ')), ' ', '-'))\
  .withColumn('end_date', F.regexp_replace(F.initcap(F.regexp_replace('end_date', '-', ' ')), ' ', '-'))

df.show()

+----+-----------+-----------+
|  id| start_date|   end_date|
+----+-----------+-----------+
|val1|01-Jan-2021|01-Mar-2021|
|val2|02-Jul-2021|01-Aug-2022|
|val3|02-Jul-2021|       null|
+----+-----------+-----------+
user_74654620
Related