Фильтровать столбец массива в кадре данных на основе заданного входного массива --Pyspark - PullRequest
Новая
       5

Фильтровать столбец массива в кадре данных на основе заданного входного массива --Pyspark

1 голос
/ 25 марта 2020

У меня есть датафрейм, подобный этому 

Studentname  Speciality
Alex         ["Physics","Math","biology"]
Sam          ["Economics","History","Math","Physics"]
Claire       ["Political science,Physics"]

Я хочу найти всех студентов, которые специализируются по [физике, математике], поэтому на выходе должно быть 2 строки Alex, Sam

Это то, что я пробовал 

from pyspark.sql.functions import array_contains
from pyspark.sql import functions as F

def student_info():
     student_df = spark.read.parquet("s3a://studentdata")
     a1=["Physics","Math"]
     df=student_df
     for a in a1:
       df= student_df.filter(array_contains(student_df.Speciality, a))
       print(df.count())

student_info()

output:
3
2

Хотелось бы узнать, как фильтровать столбец массива на основе заданного подмножества массива

Ответы [ 3 ]

2 голосов
/ 25 марта 2020

Использование функции более высокого порядка filter должно быть наиболее масштабируемым и эффективным способом сделать это ( Spark2.4 ) 

from pyspark.sql import functions as F
df.withColumn("new", F.size(F.expr("""filter(Speciality, x-> x=='Math' or x== 'Physics')""")))\
  .filter("new=2").drop("new").show(truncate=False)
+-----------+-----------------------------------+
|Studentname|Speciality                         |
+-----------+-----------------------------------+
|Alex       |[Physics, Math, biology]           |
|Sam        |[Economics, History, Math, Physics]|
+-----------+-----------------------------------+

Если вы хотите использовать array, например a1 - , динамически , то вы можете использовать F.array_except и F.array, а затем filter на size ( Искра 2,4 ): 

a1=['Math','Physics']
df.withColumn("array", F.array_except("Speciality",F.array(*(F.lit(x) for x in a1))))\
  .filter("size(array)= size(Speciality)-2").drop("array").show(truncate=False)

+-----------+-----------------------------------+
|Studentname|Speciality                         |
+-----------+-----------------------------------+
|Alex       |[Physics, Math, biology]           |
|Sam        |[Economics, History, Math, Physics]|
+-----------+-----------------------------------+

Чтобы получить счет , вы можете поставить .count() вместо .show()

2 голосов
/ 25 марта 2020

Здесь другой подход, использующий array_sort и оператор равенства Spark, который обрабатывает массивы как любой другой тип с предварительным условием их сортировки: 

from pyspark.sql.functions import lit, array, array_sort, array_intersect

target_ar = ["Physics", "Math"]
search_ar = array_sort(array(*[lit(e) for e in target_ar]))

df.where(array_sort(array_intersect(df["Speciality"], search_ar)) == search_ar) \
  .show(10, False)

# +-----------+-----------------------------------+
# |Studentname|Speciality                         |
# +-----------+-----------------------------------+
# |Alex       |[Physics, Math, biology]           |
# |Sam        |[Economics, History, Math, Physics]|
# +-----------+-----------------------------------+

Сначала мы находим общие элементы с array_intersect(df["Speciality"], search_ar), затем мы используем == для сравнения отсортированных массивов.

0 голосов
/ 25 марта 2020

При условии, что у вас нет дубликатов в Speciality для студента (например, 

StudentName   Speciality
SomeStudent   ['Physics', 'Math', 'Biology', 'Physics']

Вы можете использовать explode с groupby в pandas

Итак, для вашей проблемы 

# df is above dataframe
# Lookup subjects
a1 = ['Physics', 'Math']

gdata = df.explode('Speciality').groupby(['Speciality']).size().to_frame('Count')

gdata.loc[a1, 'Count']

#             Count
# Speciality
# Physics         3
# Math            2
...