將資料框行與逐行模式條件匹配并添加陣列列

假設您有一個資料框，并且您想通過添加新pattern_name列來過濾掉逐行模式 。列
的型別pattern_name應該是一個陣列，因為每一行都可能匹配多個模式。

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

由于輸入非常大，您如何以有效的方式（無 udf）執行此操作？

過去，我的 df 每行最多只有一個匹配項，所以我使用了when函式（下面的示例）。但是，如果每行有多個匹配項，則這不起作用，您需要一個陣列。

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

uj5u.com熱心網友回復：

您可以將patterns串列定義為：

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"),
]

然后使用array具有串列理解的函式，when您可以獲得模式名稱的串列列：

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         |[]                             |
# ---- ---- ------------- ------------------------------- 

您還可以patterns_df從上面的串列中創建一個資料框，然后使用 join 后跟 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")
)

標籤：

上一篇：使用Spark.SQL編碼的ApacheSpark結構化流

下一篇：訂購PysparkWindow時缺少資料