Пересечь каждую строку pyspark DataFrame, который является списком строк с основным списком строк? - PullRequest
Новая
       55

Пересечь каждую строку pyspark DataFrame, который является списком строк с основным списком строк?

0 голосов
/ 05 декабря 2018

Скажем, у меня есть DataFrame, как это.

[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'])]

И список главных слов, подобный этому: 

master_word_list = ['b', 'c']

Есть ли удобный способ фильтрации word_list по master_word_list, чтобы результирующий фрейм данных pyspark выглядел следующим образом.(Под гладким я подразумеваю без использования UDF, если бы UDF были лучшим / единственным способом, я бы также принял это как решение) 

[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'])]

1 Ответ

0 голосов
/ 05 декабря 2018

array_intersect доступно с версии Spark 2.4:

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

Функция сбора: возвращает массив элементов на пересечении col1 иcol2, без дубликатов.

Параметры:

  • col1 - имя столбца, содержащего массив
  • col2 - имя столбца, содержащего массив 
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]|
+-----------+---------+
...