Match dataframe rows with row-wise pattern conditions and add array column

user_3955022 0
Viewed 105

Suppose you have a dataframe, and you want to filter out row-wise patterns by adding a new pattern_name column.
The type of the pattern_name column should be an array, because each row can potentially match to multiple patterns.

# Input
df = spark.createDataFrame(
    [(1, 21, 'A foo text'), 
     (2, 42, 'A foo'),
     (2, 42, 'A foobar text'),
     (2, 42, 'barz'),],
     ['id_1', 'id_2', 'text']
)
# Patterns:
# pattern_foo_1:  id_1 = 1, id_2 = 21, text.rlike('foo')
# pattern_foo_2:  id_1 = 2, id_2 = 42, text.rlike('foo')
# pattern_foobar: id_1 = 2, id_2 = 42, text.rlike('foobar')

# Desired output: (null can also be empty string, doesn't matter)
+------+------+----------------+------------------------------------+
|  id_1|  id_2|            text|                        pattern_name|
+------+------+----------------+------------------------------------+
|     1|    21|    'A foo text'|                 ['pattern_foo_1', ]|
|     2|    42|         'A foo'|                 ['pattern_foo_2', ]|
|     2|    42| 'A foobar text'| ['pattern_foo_2', 'pattern_foobar']|
|     2|    42|          'barz'|                                null|
+------+------+----------------+------------------------------------+

How do you do this in an efficient way (no udf), because the input is very large?

In the past, my df only had a single match per row at max, so I used the when function (example below). But this does not work if you have multiple matches per row, where you'd need an array.

pattern_name_col = None
for pattern in pattern_list:
    if pattern_name_col is None:
        # pseudocode example
        pattern_name_col = when(
                                (col('id_1') == 1) & (col('id_2') == 21)
                                & (col('text').rlike('foo')),
                                'pattern_foo_1')
    else:
        pattern_name_col = pattern_name_col.when(..., ...)

df = df.withColumn('pattern_name', pattern_name_col).filter(col('pattern_name').isNotNull())
pyspark apache-spark-sql apache-spark python
1 Answers

You can define your patterns list as:

patterns = [
    (1, 21, "foo", "pattern_foo_1"), # (id_1, id_2, pattern, pattern_name)
    (2, 42, "foo", "pattern_foo_2"),
    (2, 42, "foobar", "pattern_foobar"),
]

Then using array function with list comprehension and when you can get your list column of pattern names:

import pyspark.sql.functions as F

df1 = df.withColumn(
    "pattern_name",
    F.array(*[
        F.when((F.col("id_1") == p[0]) & (F.col("id_2") == p[1]) & F.col("text").rlike(p[2]), p[3])
        for p in patterns
    ])
).withColumn(
    "pattern_name",
    F.expr("filter(pattern_name, x -> x is not null)")
)

df1.show(truncate=False)
#+----+----+-------------+-------------------------------+
#|id_1|id_2|text         |pattern_name                   |
#+----+----+-------------+-------------------------------+
#|1   |21  |A foo text   |[pattern_foo_1]                |
#|2   |42  |A foo        |[pattern_foo_2]                |
#|2   |42  |A foobar text|[pattern_foo_2, pattern_foobar]|
#|2   |42  |barz         |[]                             |
#+----+----+-------------+-------------------------------+

You can also create a patterns_df dataframe from the above list then use join followed by goupby + collect_list:

patterns_df = spark.createDataFrame(patterns, ["id_1", "id_2", "pattern", "pattern_name"])

df1 = df.alias("df").join(
    patterns_df.alias("p"),
    F.expr("df.id_1 = p.id_1 and df.id_2 = p.id_2 and df.text rlike p.pattern")
).groupBy("df.id_1", "df.id_2", "text").agg(
    F.collect_list("pattern_name").alias("pattern_name")
)
user_13865510
Related