根據是否包含substring過濾Pyspark Dataframe列

Question

我有一個 pyspark dataframe message_df有數百萬行，看起來像這樣

ID	信息
ab123	你好，我叫克里斯
CD345	房間應該是2301
ef567	歡迎？ 你叫什么名字？
gh873	請往那邊走
kj893	當前年份是 2022 年

和兩個列表

wanted_words = ['name','room']
unwanted_words = ['welcome','year']

我只想獲取message包含wanted_words中的任何單詞且不包含unwanted_words中的任何單詞的行，因此結果應該是：

ID	信息
ab123	你好，我叫克里斯
CD345	房間應該是2301

截至目前，我正在逐字逐句地進行

message_df.select(lower(F.col('message'))).filter(
    (
        F.col('lower(message)').contains('name') |
        F.col('lower(message)').contains('room') 
    ) & (
        ~F.col('lower(message)').contains('welcome') &
        ~F.col('lower(message)').contains('year') 
    )  
)

這對代碼來說非常乏味。 但是，當我改為使用rlike時：

wanted_words ="(name|room)"
unwanted_words ="(welcome|year)"

message_df.select(lower(F.col('message'))).filter(
   ~F.col('lower(message)').rlike(not_contain) &
    F.col('lower(message)').rlike(contain)
)

這個過程大大減慢。 原因是因為rlike明顯更慢，如果是的話，當wanted_words和unwanted_words可能包含數百個單詞時，什么是更好的過濾方法？

Answer 1

將文本拆分為標記/單詞並使用arrays_overlap function 檢查是否存在需要或不需要的標記：

df = df.filter(
    (
      F.arrays_overlap(
          F.split(F.regexp_replace(F.lower("message"), r"[^a-zA-Z0-9\s]+", ""), "\s+"),
          F.array([F.lit(c) for c in wanted_words])
          )
    )
    & 
    (
      ~F.arrays_overlap(
          F.split(F.regexp_replace(F.lower("message"), r"[^a-zA-Z0-9\s]+", ""), "\s+"),
          F.array([F.lit(c) for c in unwanted_words])
          )
    )
)

完整示例：

columns = ["id","message"]
data = [["ab123","Hello my name is Chris"],["cd345","The room should be 2301"],["ef567","Welcome! What is your name?"],["gh873","That way please"],["kj893","The current year is 2022"]]
df = spark.createDataFrame(data).toDF(*columns)

wanted_words = ['name','room']
unwanted_words = ['welcome','year']

df = df.filter(
    (
      F.arrays_overlap(
          F.split(F.regexp_replace(F.lower("message"), r"[^a-zA-Z0-9\s]+", ""), "\s+"),
          F.array([F.lit(c) for c in wanted_words])
          )
    )
    & 
    (
      ~F.arrays_overlap(
          F.split(F.regexp_replace(F.lower("message"), r"[^a-zA-Z0-9\s]+", ""), "\s+"),
          F.array([F.lit(c) for c in unwanted_words])
          )
    )
)

[Out]:
+-----+------------------------+
|id   |message                 |
+-----+------------------------+
|ab123|Hello my name is Chris  |
|cd345|The room should be 2301 |
+-----+------------------------+

您還可以一次預先計算令牌以提高效率：

df = df.withColumn("tokens", F.split(F.regexp_replace(F.lower("message"), r"[^a-zA-Z0-9\s]+", ""), "\s+"))

並在“arrays_overlap”中使用：

F.arrays_overlap(F.col("tokens"), F.array([F.lit(c) for c in wanted_words]))