PySpark Groupby и фильтр на основе регулярных выражений

Question

У меня есть PySpark df с

from pyspark.sql import functions as F
print(df.groupBy(['issue_month', 'loan_status']).count().show())

+-----------+------------------+-----+
|issue_month|       loan_status|count|
+-----------+------------------+-----+
|         06|        Fully Paid|12632|
|         03|        Fully Paid|16243|
|         07|           Default|    1|
|         02|        Fully Paid|16467|
|         06|           Default|    1|
|         07|   In Grace Period|  289|
|         01|       Charged Off| 5975|
|         05|       Charged Off| 5209|
|         02|Late (31-120 days)|  184|
|         11|           Current|17525|
|         12|   In Grace Period|  369|
|         10|        Fully Paid|19222|
|         04|        Fully Paid|16802|
|         07|       Charged Off| 7072|
|         06|       Charged Off| 4589|
|         04| Late (16-30 days)|   98|
|       null|              null|    2|
|         10|Late (31-120 days)|  621|
|         07| Late (16-30 days)|  125|
|         10|           Default|    2|
+-----------+------------------+-----+

Я бы хотел отфильтровать только поздний loan_status, который может принимать значения «Поздно (16-30 дней)» или «Поздно (31-120 дней) ». Поэтому я попытался:

print(df.groupBy(['issue_month', 'loan_status']).count().filter((F.col('loan_status')=='Late (31-120 days)')|F.col('loan_status')=='Late (16-30 days)').show())

Это не удалось, но, тем не менее, это грязно. Я хотел бы сделать, как в pandas, где я могу просто фильтровать регулярное выражение. В моем случае это было бы что-то вроде:

F.col('loan_status').contains("Late")

Answer 1

Pyspark также имеет функцию contains() (или) like, которую мы можем использовать в .filter()

Example:

#sample data
df.show()
#+-----------+------------------+
#|issue_month|       loan_status|
#+-----------+------------------+
#|         10|        Fully Paid|
#|         10|           Default|
#|         10|Late (31-120 days)|
#+-----------+------------------+

#in filter query convert loan_status to lower case and look for substring late.
df.groupBy("issue_month","loan_status").\
count().\
filter(lower(col("loan_status")).contains("late")).\
show()

#by using like function
df.groupBy("issue_month","loan_status").\
count().\
filter(lower(col("loan_status")).like("late%")).\
show()

#i would suggest filtering rows before groupby will significantly increases the performance in bigdata!!
df.filter(lower(col("loan_status")).like("late%")).\
groupBy("issue_month","loan_status").\
count().\
show()

#+-----------+------------------+-----+
#|issue_month|       loan_status|count|
#+-----------+------------------+-----+
#|         10|Late (31-120 days)|    1|
#+-----------+------------------+-----+

---

Мы можем использовать .agg(sum("count")), чтобы получить сумму подсчета независимо от номера_месяца.

Example:

from pyspark.sql.functions import sum as _sum
df.show()
#+-----------+------------------+
#|issue_month|       loan_status|
#+-----------+------------------+
#|         10|        Fully Paid|
#|         10|           Default|
#|         11|Late (31-120 days)|
#|         11|Late (31-120 days)|
#|         10| Late (16-30 days)|
#+-----------+------------------+

df.filter(lower(col("loan_status")).contains("late")).\
groupBy("issue_month","loan_status").\
count().\
agg(_sum("count").alias("sum")).\
show()

#+---+
#|sum|
#+---+
#|  3|
#+---+

---

df.filter(lower(col("loan_status")).like("late%")).\
groupBy("issue_month","loan_status").\
count().\
groupBy("loan_status").\
agg(_sum("count").alias("sum_count")).\
show()

#same result will get by using one group too
df.filter(lower(col("loan_status")).contains("late")).\
groupBy("loan_status").\
agg(count("*").alias("sum_count")).\
show()

#+------------------+---------+
#|       loan_status|sum_count|
#+------------------+---------+
#|Late (31-120 days)|        2|
#| Late (16-30 days)|        1|
#+------------------+---------+

---

ОБНОВЛЕНИЕ:

df.filter(lower(col("loan_status")).contains("late")).\
groupBy("issue_month").\
agg(count("*").alias("sum_count")).\
show()

#+-----------+---------+
#|issue_month|sum_count|
#+-----------+---------+
#|         10|        1|
#|         11|        2|
#+-----------+---------+