Заменить определенные поля в данных в зависимости от условий

Question

У меня есть датафрейм следующим образом. Следующее предназначено только для 1 пациента и 1 конкретного теста. Он может иметь несколько других тестов с похожим видом.

ptid,blast_date,test_name,result_date,test_result,date_diff
PT381201021,2019-08-22,Albumin,2019-08-14,4.3,8
PT381201021,2019-05-17,Albumin,NA,NA,0
PT381201021,2019-05-18,Albumin,NA,NA,0
PT381201021,2019-05-21,Albumin,NA,NA,0
PT381201021,2019-05-23,Albumin,NA,NA,0
PT381201021,2019-05-16,Albumin,NA,NA,0
PT381201021,2019-05-19,Albumin,NA,NA,0
PT381201021,2019-05-22,Albumin,NA,NA,0
PT381201021,2019-05-20,Albumin,NA,NA,0

Я хочу, чтобы result_date, test_result для "Albumin" в этом случае был заполнен из предыдущего blast_date, если он находится ниже определенного порога, давайте предположим, что в этом случае 3 месяца. Поэтому я хочу, чтобы следующая строка была заполнена следующим образом:

PT381201021,2019-05-23,Albumin,2019-08-14,4.3,0

Вы можете оставить столбец date_diff без изменений.

Итак, окончательный кадр данных ожидается следующим образом: -

ptid,blast_date,test_name,result_date,test_result,date_diff
PT381201021,2019-08-22,Albumin,2019-08-14,4.3,8
PT381201021,2019-05-17,Albumin,NA,NA,0
PT381201021,2019-05-18,Albumin,NA,NA,0
PT381201021,2019-05-21,Albumin,NA,NA,0
PT381201021,2019-05-23,Albumin,2019-08-14,4.3,0
PT381201021,2019-05-16,Albumin,NA,NA,0
PT381201021,2019-05-19,Albumin,NA,NA,0
PT381201021,2019-05-22,Albumin,NA,NA,0
PT381201021,2019-05-20,Albumin,NA,NA,0

Я пытался использовать функцию задержки, но у меня возникли некоторые трудности. В поисках pyspark способа решить эту проблему.

Answer 1

Вы должны использовать window functions, с rangeBetween на seconds.

from pyspark.sql import functions as F
from pyspark.sql.window import Window

w=Window().partitionBy("ptid","test_name").orderBy(F.to_timestamp("blast_date","yyyy-MM-dd").cast("long")).rangeBetween(Window.currentRow,86400*91)

df.withColumn("collect", F.collect_list(F.array("result_date","test_result")).over(w))\
  .withColumn("collect", F.expr("""filter(collect,x-> array_contains(x,'NA')!=True)""")[0])\
  .withColumn("result_date", F.when((F.col("result_date")=='NA')&(F.col("collect").isNotNull()),F.col("collect")[0]).otherwise(F.col("result_date")))\
  .withColumn("test_result", F.when((F.col("test_result")=='NA')&(F.col("collect").isNotNull()),F.col("collect")[1]).otherwise(F.col("test_result"))).drop("timestamp","collect").show(truncate=False)

+-----------+----------+---------+-----------+-----------+---------+
|ptid       |blast_date|test_name|result_date|test_result|date_diff|
+-----------+----------+---------+-----------+-----------+---------+
|PT381201021|2019-05-16|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-17|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-18|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-19|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-20|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-21|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-22|Albumin  |NA         |NA         |0        |
|PT381201021|2019-05-23|Albumin  |2019-08-14 |4.3        |0        |
|PT381201021|2019-08-22|Albumin  |2019-08-14 |4.3        |8        |
+-----------+----------+---------+-----------+-----------+---------+

Answer 2

Надеюсь, что этот подход поможет, хотя и не очень оптимизированный и с потоком выполнения, он может быть дополнительно оптимизирован.

df = spark.read.csv("/Users/61471871.csv", header=True, inferSchema=True)
df2 = df.withColumn("start_date",     F.to_date(df.blast_date)).withColumn("end_date",     F.add_months(F.to_date(df.blast_date),3)).sort(df.start_date.desc())
df_right = df2.sort(df.blast_date.desc())
df3.createOrReplaceTempView("tbl")
spark.sql("select * from tbl").show()
'''
|       ptid|         blast_date|test_name|result_date|test_result|date_diff|
+-----------+-------------------+---------+-----------+-----------+---------+
|PT381201021|2019-08-22 00:00:00|  Albumin| 2019-08-14|        4.3|        8|
|PT381201021|2019-05-23 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-22 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-21 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-20 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-19 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-18 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-17 00:00:00|  Albumin|         NA|         NA|        0|
|PT381201021|2019-05-16 00:00:00|  Albumin|         NA|         NA|        0|
+-----------+-------------------+---------+-----------+-----------+---------+
'''
df2 = df.sort(df.blast_date.desc).withColumn("90_days_back" ,F.add_months(to_date(df.blast_date), 3)).show()
df2 = df.select(F.add_months(df.blast_date, 3).alias('third_month'))


df_left = spark.sql("select ptid, max(start_date) as range_dt from tbl group by ptid ")
df_one = df_right.crossJoin(df_left)

df_right.join(df_left, df_left.ptid = df_right.ptid).show()
df_two = df_one.withColumn("date_diff", F.datediff(df_one.start_date,     df_one.range_dt))
'''
+-----------+-------------------+---------+-----------+-----------+---------+----------+----------+-----------+----------+
|       ptid|         blast_date|test_name|result_date|test_result|date_diff|start_date|  end_date|       ptid|  range_dt|
+-----------+-------------------+---------+-----------+-----------+---------+----------+----------+-----------+----------+
|PT381201021|2019-08-22 00:00:00|  Albumin| 2019-08-14|        4.3|        0|2019-08-22|2019-11-22|PT381201021|2019-08-22|
|PT381201021|2019-05-23 00:00:00|  Albumin|         NA|         NA|      -91|2019-05-23|2019-08-23|PT381201021|2019-08-22|
|PT381201021|2019-05-22 00:00:00|  Albumin|         NA|         NA|      -92|2019-05-22|2019-08-22|PT381201021|2019-08-22|
|PT381201021|2019-05-21 00:00:00|  Albumin|         NA|         NA|      -93|2019-05-21|2019-08-21|PT381201021|2019-08-22|
|PT381201021|2019-05-20 00:00:00|  Albumin|         NA|         NA|      -94|2019-05-20|2019-08-20|PT381201021|2019-08-22|
|PT381201021|2019-05-19 00:00:00|  Albumin|         NA|         NA|      -95|2019-05-19|2019-08-19|PT381201021|2019-08-22|
|PT381201021|2019-05-18 00:00:00|  Albumin|         NA|         NA|      -96|2019-05-18|2019-08-18|PT381201021|2019-08-22|
|PT381201021|2019-05-17 00:00:00|  Albumin|         NA|         NA|      -97|2019-05-17|2019-08-17|PT381201021|2019-08-22|
|PT381201021|2019-05-16 00:00:00|  Albumin|         NA|         NA|      -98|2019-05-16|2019-08-16|PT381201021|2019-08-22|
+-----------+-------------------+---------+-----------+-----------+---------+----------+----------+-----------+----------+
'''

Теперь у вас есть флаг разницы дат, вы можете применить фильтр и затем сделать соединение, чтобы получить ожидаемый результат.

Код может быть дополнительно оптимизирован для работы с большим набором данных.