Как фильтровать значения столбца в массиве в Pyspark - PullRequest
Новая
       12

Как фильтровать значения столбца в массиве в Pyspark

0 голосов
/ 08 ноября 2019

У меня есть pyspark Dataframe, который содержит много столбцов, в том числе столбец типа Array и столбец String: 

numbers  <Array>              |    name<String>
------------------------------|----------------
["160001","160021"]           |     A
------------------------------|----------------
["160001","1600", "42345"]    |     B
------------------------------|----------------
["160001","9867", "42345"]    |     C
------------------------------|----------------
["160001","8650", "2345"]     |     A
------------------------------|----------------
["2456","78568", "42345"]     |     B
-----------------------------------------------

Я хочу пропустить числа, содержащие 4 цифры, из столбца чисел if the name column is not "B". And keep it if the name column is "B". Например:

In the lines 2 and 5, I have "1600" and "2456" contains 4 digits и столбец имени «B», я должен хранить их в значениях столбца: 

------------------------------|----------------
["160001","1600", "42345"]    |     B
------------------------------|----------------
["2456","78568", "42345"]     |     B
-----------------------------------------------

В строках 3 и 4, У меня есть столбец чисел, который содержит цифры 4 цифры, но имя столбца отличается от "B" ==> Поэтому я должен пропустить их.

Пример: 

------------------------------|----------------
["160001","9867", "42345"]    |     C
------------------------------|----------------
["160001","8650", "2345"]     |     A
------------------------------|----------------

Ожидаемый результат: 

    numbers  <Array>              |    name<String>
------------------------------|----------------
["160001","160021"]           |     A
------------------------------|----------------
["160001","1600", "42345"]    |     B
------------------------------|----------------
["160001", "42345"]           |     C
------------------------------|----------------
["160001"]                    |     A
------------------------------|----------------
["2456","78568", "42345"]     |     B
-----------------------------------------------

Как я могу это сделать? Спасибо

Ответы [ 2 ]

2 голосов
/ 08 ноября 2019

Начиная с Spark 2.4, вы можете использовать функцию более высокого порядка FILTER для фильтрации массива. Объединение этого с выражением if должно решить проблему: 

df.selectExpr("if(name != \'B', FILTER(numbers, x -> length(x) != 4), numbers) AS numbers", "name")
1 голос
/ 08 ноября 2019

Вам нужно написать udf для фильтра массива и использовать его с предложением when, чтобы применить udf к определенному условию, например where name == B: 

from pysparl.sql.functions import udf, col, when
from pyspark.sql.types import ArrayType, StringType

filter_array_udf = udf(lambda arr: [x for x in arr if len(x) > 4], "array<string>")
# if string schema is not supported then use the next commented line
# filter_array_udf = udf(filter_array, ArrayType(StringType()))

df = df.withColumn("numbers", when(col("name") == "B", filter_array_udf(col("numbers"))).otherwise(col("numbers")))
...