UDF with multiple rows as response pySpark - apache-spark

Question

I want to apply splitUtlisation on each row of utilisationDataFarme and pass startTime and endTime as parameters., as a result splitUtlisation will return multiple rows of data hence I want to create a new DataFrame with (Id, Day, Hour, Minute)
def splitUtlisation(onDateTime, offDateTime):
yield onDateTime
rule = rrule.rrule(rrule.HOURLY, byminute = 0, bysecond = 0, dtstart=offDateTime)
for result in rule.between(onDateTime, offDateTime):
yield result
yield offDateTime
utilisationDataFarme = (
sc.parallelize([
(10001, "2017-02-12 12:01:40" , "2017-02-12 12:56:32"),
(10001, "2017-02-13 12:06:32" , "2017-02-15 16:06:32"),
(10001, "2017-02-16 21:45:56" , "2017-02-21 21:45:56"),
(10001, "2017-02-21 22:32:41" , "2017-02-25 00:52:50"),
]).toDF(["id", "startTime" , "endTime"])
.withColumn("startTime", col("startTime").cast("timestamp"))
.withColumn("endTime", col("endTime").cast("timestamp"))
In core Python I did like this
dayList = ['SUN' , 'MON' , 'TUE' , 'WED' , 'THR' , 'FRI' , 'SAT']
for result in hours_aligned(datetime.datetime.now(), datetime.datetime.now() + timedelta(hours=68)):
print(dayList[datetime.datetime.weekday(result)], result.hour, 60 if result.minute == 0 else result.minute)
Result
THR 21 60
THR 22 60
THR 23 60
FRI 0 60
FRI 1 60
FRI 2 60
FRI 3 60
How to create it in pySpark?
I tried to create new Schema and apply
schema = StructType([StructField("Id", StringType(), False), StructField("Day", StringType(), False), StructField("Hour", StringType(), False) , StructField("Minute", StringType(), False)])
udf_splitUtlisation = udf(splitUtlisation, schema)
df = sqlContext.createDataFrame([],"id" , "Day" , "Hour" , "Minute")
Still I could not handle multiple rows as response.

Answer 1

You can use pyspark's explode to unpack a single row containing multiple values into multiple rows once you have your udf defined correctly.
As far as I know you won't be able to use generators with yield as an udf. Instead, you need to return all values at once as an array (see return_type) which then can be exploded and expanded:
import pandas as pd
from pyspark.sql.functions import col, udf, explode
from pyspark.sql.types import ArrayType, StringType, MapType
# input data as given by OP
df = (
sc.parallelize(
[
(10001, "2017-02-12 12:01:40", "2017-02-12 12:56:32"),
(10001, "2017-02-13 12:06:32", "2017-02-15 16:06:32"),
(10001, "2017-02-16 21:45:56", "2017-02-21 21:45:56"),
(10001, "2017-02-21 22:32:41", "2017-02-25 00:52:50"),
]
)
.toDF(["id", "startTime", "endTime"])
.withColumn("startTime", col("startTime").cast("timestamp"))
.withColumn("endTime", col("endTime").cast("timestamp"))
)
return_type = ArrayType(MapType(StringType(), StringType()))
#udf(returnType=return_type)
def your_udf_func(start, end):
"""Insert your function to return whatever you like
as a list of dictionaries.
For example, I chose to return hourly values for
day, hour and minute.
"""
date_range = pd.date_range(start, end, freq="h")
df = pd.DataFrame(
{
"day": date_range.strftime("%a"),
"hour": date_range.hour,
"minute": date_range.minute,
}
)
values = df.to_dict("index").values()
return list(values)
extracted = your_udf_func("startTime", "endTime")
exploded = explode(extracted).alias("exploded")
expanded = [
col("exploded").getItem(k).alias(k) for k in ["hour", "day", "minute"]
]
result = df.select("id", exploded).select("id", *expanded)
And the result is:
result.show(5)
+-----+----+---+------+
| id|hour|day|minute|
+-----+----+---+------+
|10001| 12|Sun| 1|
|10001| 12|Mon| 6|
|10001| 13|Mon| 6|
|10001| 14|Mon| 6|
|10001| 15|Mon| 6|
+-----+----+---+------+
only showing top 5 rows