简体   繁体   中英

Intersect each row of a pyspark DataFrame which is a list of strings with a master list of strings?

Say I have a DataFrame like this.

[Row(case_number='5307793179', word_list=['n', 'b', 'c']),
 Row(case_number='5307793171', word_list=['w', 'e', 'c']),
 Row(case_number='5307793172', word_list=['1', 'f', 'c']),
 Row(case_number='5307793173', word_list=['a', 'k', 'c']),
 Row(case_number='5307793174', word_list=['z', 'l', 'c']),
 Row(case_number='5307793175', word_list=['b', 'r', 'c'])]

And a master word list like this:

master_word_list = ['b', 'c']

Is there a sleek way to filter word_list against master_word_list so the resulting pyspark dataframe looks like this. (by sleek I mean without using UDFs, if UDFs are the best/only way, I'd accept that as a solution as well)

[Row(case_number='5307793179', word_list=['b', 'c']),
 Row(case_number='5307793171', word_list=['c']),
 Row(case_number='5307793172', word_list=['c']),
 Row(case_number='5307793173', word_list=['c']),
 Row(case_number='5307793174', word_list=['c']),
 Row(case_number='5307793175', word_list=['b', 'c'])]

array_intersect available since Spark 2.4:

pyspark.sql.functions.array_intersect(col1, col2)

Collection function: returns an array of the elements in the intersection of col1 and col2, without duplicates.

Parameters:

  • col1 – name of column containing array
  • col2 – name of column containing array
from pyspark.sql.functions import array, array_intersect, lit

master_word_list_col = array(*[lit(x) for x in master_word_list])

df = spark.createDataFrame(
    [("5307793179", ["n", "b", "c"])], 
    ("case_number", "word_list")
)

df.withColumn("word_list", array_intersect("word_list", master_word_list_col)).show()
+-----------+---------+
|case_number|word_list|
+-----------+---------+
| 5307793179|   [b, c]|
+-----------+---------+

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM