assign patient to nearest facility in spark,scala - scala

Question

I have a data frame as below
+-----------+----------+-----------+--------+-------------+
|LOCATION_ID|PATIENT_ID|FACILITY_ID|DISTANCE|rank_distance|
+-----------+----------+-----------+--------+-------------+
|LOC0001 |P1 |FAC003 |54 |2 |
|LOC0001 |P1 |FAC002 |45 |1 |
|LOC0001 |P2 |FAC003 |54 |2 |
|LOC0001 |P2 |FAC002 |45 |1 |
|LOC0010 |P3 |FAC006 |12 |1 |
|LOC0010 |P3 |FAC003 |54 |4 |
|LOC0010 |P3 |FAC005 |23 |2 |
|LOC0010 |P3 |FAC002 |45 |3 |
|LOC0010 |P4 |FAC002 |45 |3 |
|LOC0010 |P4 |FAC005 |23 |2 |
|LOC0010 |P4 |FAC003 |54 |4 |
|LOC0010 |P4 |FAC006 |12 |1 |
|LOC0010 |P5 |FAC006 |12 |1 |
|LOC0010 |P5 |FAC002 |45 |3 |
|LOC0010 |P5 |FAC005 |23 |2 |
|LOC0010 |P5 |FAC003 |54 |4 |
|LOC0010 |P6 |FAC006 |12 |1 |
|LOC0010 |P6 |FAC005 |23 |2 |
|LOC0010 |P6 |FAC002 |45 |3 |
|LOC0010 |P6 |FAC003 |54 |4 |
|LOC0043 |P7 |FAC004 |42 |1 |
|LOC0054 |P8 |FAC002 |24 |2 |
|LOC0054 |P8 |FAC006 |12 |1 |
|LOC0054 |P8 |FAC005 |76 |3 |
|LOC0054 |P8 |FAC007 |100 |4 |
|LOC0065 |P9 |FAC006 |32 |1 |
|LOC0065 |P9 |FAC005 |54 |2 |
|LOC0065 |P10 |FAC006 |32 |1 |
|LOC0065 |P10 |FAC005 |54 |2 |
+-----------+----------+-----------+--------+-------------+
for each patient I have to assign facility for which rank is least.my output map should be as below
p1 ---> FAC002 (because its rank is least)
P2 ---> FAC002 (because its rank is least)
note each facility has capacity of just 2,except of FAC003 which has capacity of 3
so for P3,P4,P5 and P6 output should be
p3 ----> FAC006 (because its rank is 1)
P4 ----> FAC006 (because its rank is 1)
p5 ----> FAC005 (bacause FAC006 has fulled its capacity of 2,and now least
rank is of FAC005)
p6 ---->FAC005 (bacause FAC005 has one capacity left)
P7 ----->FAC004

Answer 1

Please check below code.
Extracting data from DataFrame based on patients & checking condition with previously extracted data, if the condition is false skipping that record else join that records with previous extracted records.
scala> :paste
// Entering paste mode (ctrl-D to finish)
spark.time {
import spark.implicits._
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions._
val df = Seq(("LOC0001","P1","FAC003","54","2"),("LOC0001","P1","FAC002","45","1"),("LOC0001","P2","FAC003","54","2"),("LOC0001","P2","FAC002","45","1"),("LOC0010","P3","FAC006","12","1"),("LOC0010","P3","FAC003","54","4"),("LOC0010","P3","FAC005","23","2"),("LOC0010","P3","FAC002","45","3"),("LOC0010","P4","FAC002","45","3"),("LOC0010","P4","FAC005","23","2"),("LOC0010","P4","FAC003","54","4"),("LOC0010","P4","FAC006","12","1"),("LOC0010","P5","FAC006","12","1"),("LOC0010","P5","FAC002","45","3"),("LOC0010","P5","FAC005","23","2"),("LOC0010","P5","FAC003","54","4"),("LOC0010","P6","FAC006","12","1"),("LOC0010","P6","FAC005","23","2"),("LOC0010","P6","FAC002","45","3"),("LOC0010","P6","FAC003","54","4"),("LOC0043","P7","FAC004","42","1"),("LOC0054","P8","FAC002","24","2"),("LOC0054","P8","FAC006","12","1"),("LOC0054","P8","FAC005","76","3"),("LOC0054","P8","FAC007","100","4"),("LOC0065","P9","FAC006","32","1"),("LOC0065","P9","FAC005","54","2"),("LOC0065","P10","FAC006","32","1"),("LOC0065","P10","FAC005","54","2")).toDF("location_id","patient_id","facility_id","distance","rank_distance").withColumn("facility_new",first($"facility_id").over(Window.partitionBy($"patient_id").orderBy($"rank_distance".asc))).orderBy(substring($"patient_id",2,3).cast("int").asc)
val patients = df.select("patient_id").distinct.orderBy(substring($"patient_id",2,3).cast("int").asc).collect.map(_.getAs[String](0)) // Taking all patients into collection.
val facilities = Map("FAC004" -> 2, "FAC003" -> 3, "FAC007" -> 2, "FAC002" -> 2, "FAC006" -> 2, "FAC005" -> 2)// For checking conditions.
case class Config(newDF: DataFrame,oldDF: DataFrame,facilities: Map[String,Int]) // Taking all facilitiy ids & max allowed values for checking codition.
def fetchFacilityId(config:Config) = {
config.newDF.select("facility_id").distinct.orderBy(substring($"facility_id",5,6).cast("int").asc)
.except(config.oldDF.select("facility_new").distinct.orderBy(substring($"facility_new",5,6).cast("int").asc)).head.getAs[String](0)
} // Getting required facility id.
def findFacilityId(config:Config, patientId: String,index:Int):(Boolean,String,Int) = {
val max_distance = config.oldDF.filter($"patient_id"=== patientId).select("rank_distance").orderBy($"rank_distance".desc).head.getAs[String](0).toInt
(index < max_distance) match {
case true => {
val fac = config.oldDF.filter($"patient_id" === patientId && $"rank_distance"=== index).select("facility_id").distinct.map(_.getAs[String](0)).collect.head
val bool = config.newDF.filter($"facility_new" === lit(fac) && $"rank_distance"=== index).select("patient_id").distinct.count < config.facilities(fac)
(bool,fac, index)
}
case false => (true,fetchFacilityId(config), 0)
}
} // finding required facility id.
def process(config:Config,patientId:String,index:Int):DataFrame = findFacilityId(config,patientId,index) match {
case (true,fac,_) => {
config.newDF.union(config.oldDF.filter($"patient_id" === patientId).withColumn("facility_new",lit(fac)))
}
case (false,_,rank) => {
process(config,patientId,index+1)
}
} // Checking duplicate facility ids
val config = Config(newDF= df.limit(0),df, facilities)
val updatedDF = patients.foldLeft(config){ case ((cfg),patientId) =>
cfg.newDF.count match {
case 0L => cfg.copy(newDF = cfg.newDF.union(cfg.oldDF.filter($"patient_id" === patientId)))
case _ => cfg.copy(newDF = process(cfg, patientId,1))
}
}.newDF.drop("facility_id").withColumnRenamed("facility_new","facility_id")
updatedDF.show(31, false)
// Exiting paste mode, now interpreting.
+-----------+----------+--------+-------------+-----------+
|location_id|patient_id|distance|rank_distance|facility_id|
+-----------+----------+--------+-------------+-----------+
|LOC0001 |P1 |45 |1 |FAC002 |
|LOC0001 |P1 |54 |2 |FAC002 |
|LOC0001 |P2 |54 |2 |FAC002 |
|LOC0001 |P2 |45 |1 |FAC002 |
|LOC0010 |P3 |12 |1 |FAC006 |
|LOC0010 |P3 |54 |4 |FAC006 |
|LOC0010 |P3 |23 |2 |FAC006 |
|LOC0010 |P3 |45 |3 |FAC006 |
|LOC0010 |P4 |45 |3 |FAC006 |
|LOC0010 |P4 |23 |2 |FAC006 |
|LOC0010 |P4 |54 |4 |FAC006 |
|LOC0010 |P4 |12 |1 |FAC006 |
|LOC0010 |P5 |12 |1 |FAC005 |
|LOC0010 |P5 |45 |3 |FAC005 |
|LOC0010 |P5 |23 |2 |FAC005 |
|LOC0010 |P5 |54 |4 |FAC005 |
|LOC0010 |P6 |12 |1 |FAC005 |
|LOC0010 |P6 |23 |2 |FAC005 |
|LOC0010 |P6 |45 |3 |FAC005 |
|LOC0010 |P6 |54 |4 |FAC005 |
|LOC0043 |P7 |42 |1 |FAC003 |
|LOC0054 |P8 |24 |2 |FAC003 |
|LOC0054 |P8 |12 |1 |FAC003 |
|LOC0054 |P8 |76 |3 |FAC003 |
|LOC0054 |P8 |100 |4 |FAC003 |
|LOC0065 |P9 |32 |1 |FAC003 |
|LOC0065 |P9 |54 |2 |FAC003 |
|LOC0065 |P10 |32 |1 |FAC003 |
|LOC0065 |P10 |54 |2 |FAC003 |
+-----------+----------+--------+-------------+-----------+
Time taken: 31009 ms
scala>