Extract regex that matches dataframe value

user_13192486 0
Viewed 205

I have a list of regex {"WeLove*", "Arizona.*hot", "Mahi*"} and a dataframe with certain values that might match with one of the Regex expressions from the regex list.

_c0 _c1 _c2 _c3
Arizona is hot 2020 1 Y
Arizona happens to be hot 2020 1 Y
MahiWalia 2020 1 Y
MahiSingh 2020 1 Y
MahiRandhawa 2020 1 Y
WeLovechocolate 2020 1 Y

I am trying to find out for each regex how many matches are there in the data frame?

"WeLove*"  = 1  
"Arizona.*hot" = 2
"Mahi*" = 3

By doing this :

regex = spark.read.csv("regex.csv", header=False)

regex_list = regex.select("_c0").rdd.flatMap(lambda x: x).collect()

output = df.select("*", when(col("_c0").rlike("|".join(regex_list)), "True").alias("matches"))

filteredDf = output.filter(output.matches.isNotNull())

I was able to find the rows in dataFrame that at least matched with one of the regex in the list but I am not sure how to find which one and then group the rows according to the regex they matched with?

pyspark apache-spark-sql apache-spark python
2 Answers

You can use coalesce along with when function to assign a regex_group for each row based on the regex_list values, then group by the created group column and count :

from pyspark.sql import functions as F

regex_list = ["WeLove*", "Arizona.*hot", "Mahi*"]

df1 = df.withColumn(
    "regex_group",
    F.coalesce(*[F.when(F.col("_c0").rlike(r), F.lit(r)) for r in regex_list])
).groupBy("regex_group").count()

df1.show()

#+------------+-----+
#| regex_group|count|
#+------------+-----+
#|Arizona.*hot|    2|
#|       Mahi*|    3|
#|     WeLove*|    1|
#+------------+-----+
user_13865510

You can join the regex dataframe to the df dataframe using an rlike condition, and then get the count for each regex:

import pyspark.sql.functions as F

regex = spark.read.csv("regex.csv", header=False)

result = df.alias('df').join(
    regex.alias('regex'),
    F.expr('df._c0 rlike regex._c0')
).groupBy('regex._c0').count()

result.show()
+------------+-----+
|         _c0|count|
+------------+-----+
|Arizona.*hot|    2|
|       Mahi*|    3|
|     WeLove*|    1|
+------------+-----+
user_141657300
Related