Map list of multiple substrings in PySpark

Jaol 0
Viewed 46

I have a dataframe with one column like this:

Locations
Germany:city_Berlin
France:town_Montpellier
Italy:village_Amalfi

I would like to get rid of the substrings: 'city_', 'town_', 'village_', etc.

So the output should be:

Locations
Germany:Berlin
France:tMontpellier
Italy:Amalfi

I can get rid of one of them this way: F.regexp_replace('Locations', 'city_', '')

Is there a similar way to pass several substrings to remove from the original column?

Ideally I'm looking for a one line solution, without having to create separate functions or convoluted things.

substring string dataframe pyspark apache-spark
3 Answers

I wouldnt map. Looks to me like you want to replace strings immediately to the left of : if they end with _. If so use regex. Code below

df.withColumn('new_Locations', regexp_replace('Locations', '(?<=\:)[a-z_]+','')).show(truncate=False)


+---+-----------------------+------------------+
|id |Locations              |new_Locations     |
+---+-----------------------+------------------+
|1  |Germany:city_Berlin    |Germany:Berlin    |
|2  |France:town_Montpellier|France:Montpellier|
|4  |Italy:village_Amalfi   |Italy:Amalfi      |
+---+-----------------------+------------------+
wwnde0 
F.regexp_replace('Locations', r'(?<=:).*_', '')

.* tells that you will match all characters. But it is located between (?<=:) and _.

_ is the symbol which must follow all the characters matched by .*.

(?<=:) is a syntax for "positive lookbehind". It is not a part of a match, but it ensures that right before the .*_ you must have a : symbol.

Another option - list of strings to remove

strings = ['city', 'big_city', 'town', 'big_town', 'village']

F.regexp_replace('Locations', fr"(?<=:)({'|'.join(strings)})_", '')

Full example:

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [('Germany:big_city_Berlin',),
     ('France:big_town_Montpellier',),
     ('Italy:village_Amalfi',)],
    ['Locations'])

df = df.withColumn('loc1', F.regexp_replace('Locations', r'(?<=:).*_', ''))

strings = ['city', 'big_city', 'town', 'big_town', 'village']
df = df.withColumn('loc2', F.regexp_replace('Locations', fr"(?<=:)({'|'.join(strings)})_", ''))

df.show(truncate=0)
# +---------------------------+------------------+------------------+
# |Locations                  |loc1              |loc2              |
# +---------------------------+------------------+------------------+
# |Germany:big_city_Berlin    |Germany:Berlin    |Germany:Berlin    |
# |France:big_town_Montpellier|France:Montpellier|France:Montpellier|
# |Italy:village_Amalfi       |Italy:Amalfi      |Italy:Amalfi      |
# +---------------------------+------------------+------------------+
ZygD0

Not 100% sure about your case although please find here another solution for your problem (Spark 3.x).

from pyspark.sql.functions import expr

df.withColumn("loc_map", expr("str_to_map(Locations, ',', ':')")) \
  .withColumn("Locations", expr("transform_values(loc_map, (k, v) -> element_at(split(v, '_'), size(split(v, '_'))))")) \
  .drop("loc_map") \
  .show(10, False)
  1. First convert Locations into a map using str_to_map
  2. Then iterate through the values of the map and transform each value using transform_values.
  3. Inside transform_values, we use element_at, split and size to identify and return the last item of an array separated by _
abiratsis0
Related