Forward filling in Spark

from datetime import datetime, timedelta
from faker import Faker
import json


DATE_END = datetime.now()
DATE_START = DATE_END - timedelta(days=31)
NUM_EVENTS = 10

class EventGenerator:
    """ Defines the EventGenerator """

    MIN_LIVES = 1
    MAX_LIVES = 99
    CHARACTERS = ["Mario", "Luigi", "Peach", "Toad"]

    def __init__(self, num_events, output_type, start_date, end_date, output_file=None):
        """ Initialize the EventGenerator """
        self.faker = Faker()
        self.num_events = num_events
        self.output_type = output_type
        self.output_file = output_file
        self.start_date = start_date
        self.end_date = end_date

    def _get_date_between(self, date_start, date_end):
        """ Get a date between start and end date """
        return self.faker.date_between_dates(date_start=date_start, date_end=date_end)

    def _generate_events(self):
        """ Generate the metric data """
        for _ in range(self.num_events):
            yield {
                "character": self.faker.random_element(self.CHARACTERS),
                "world": self.faker.random_int(min=1, max=8, step=1),
                "level": self.faker.random_int(min=1, max=4, step=1),
                "lives": self.faker.random_int(
                    min=self.MIN_LIVES, max=self.MAX_LIVES, step=1
                ),
                "time": str(self._get_date_between(self.start_date, self.end_date)),
            }

    def store_events(self):
        if self.output_type == "jl":
            with open(self.output_file, "w") as outputfile:
                for event in self._generate_events():
                    outputfile.write(f"{json.dumps(event)}\n")
        elif self.output_type == "list":
            return list(self._generate_events())

params = {
    "num_events": NUM_EVENTS,
    "output_type": "list",
    "start_date": DATE_START,
    "end_date": DATE_END,
}
# Create the event generator
generator = EventGenerator(**params)
# Create and store the events
events = generator.store_events()

import pandas as pd
pd.DataFrame(events).sort_values(["character", "time"])

from pyspark.sql import SparkSession, SQLContext
from pyspark.sql import functions as F
from pyspark.sql.window import Window

import os
os.environ['PYSPARK_PYTHON'] = "/data/jupyter/bin/python"

spark = (SparkSession
  .builder
  .appName("Spark Forward Fill")
  .getOrCreate())

sc = SQLContext(spark)

df = sc.createDataFrame(events)

df.show()

df.dtypes

df = df.withColumn("time", F.to_date(F.to_timestamp("time")))

df.dtypes

df = df.withColumn("status", F.concat(F.col('world'), F.lit('-'), F.col('level')))
df.show()

w = Window().partitionBy("character").orderBy("time")

df = df.withColumn("next_status", F.lead("time").over(w))
df.show()

df = df.withColumn("sequence", F.when(F.col("next_status").isNotNull(),
                               F.expr("sequence(to_date(time), date_sub(to_date(next_status),1), interval 1 day)"))\
                                .otherwise(F.array("time")))
df.show()

df.select("character", "status", F.explode("sequence").alias("time")).show()

# Create the dataframe
df = sc.createDataFrame(events)
# Modify the data type
df = df.withColumn("time",
                   F.to_date(F.to_timestamp("time")))
# Create the status column
df = df.withColumn("status",
                   F.concat(F.col('world'),
                            F.lit('-'),
                            F.col('level')))
# Apply the window
df = df.withColumn("next_status",
                   F.lead("time").over(w))
# Fill in the empty `next_status` column with today's date
df = df.withColumn("next_status",
                   F.when(F.col("next_status").isNull(),
                          F.expr("current_date()"))\
                           .otherwise(F.col("next_status")))
# Apply the sequence
df = df.withColumn("sequence",
                   F.when(F.col("next_status").isNotNull(),
                          F.expr("sequence(to_date(time), date_sub(to_date(next_status), 1), interval 1 day)"))\
                           .otherwise(F.array("time")))
# Select the columns and explore the sequence
df.select("character", "status", F.explode("sequence").alias("time")).show()

pdf = pd.DataFrame(events)
# Convert datatypes
pdf['time'] = pd.to_datetime(pdf['time'])
# Add status column
pdf['status'] = pdf.apply(lambda x: f"{x['world']}-{x['level']}", axis=1)
# Create full time range as frame
timeframe = pd.date_range(start=min(pdf['time']),
                          end=datetime.now().date()).to_frame().reset_index(drop=True).rename(columns={0: 'time'})
# Merge timeframe into original frame
pdf = pdf.merge(timeframe,
                left_on='time',
                right_on='time',
                how='right')
# 1. Pivot to get dates on rows and characters as columns
# 2. Forward fill values per character
# 3. Fill remaining NaNs with False
pdf = pdf.pivot(index='time',
                columns='character',
                values='status')

pdf = pdf.fillna(method='ffill')
# Drop NaN column and reset the index
pdf = pdf.loc[:, pdf.columns.notnull()].reset_index()
# Melt the columns back
pdf = pd.melt(pdf,
              id_vars='time',
              value_name='status')
print(f"Original length: {len(events)}, new length: {len(pdf)}")
pdf.head(10)

# Create the dataframe
df = sc.createDataFrame(events)
# Modify the data type
df = df.withColumn("time",
                   F.to_date(F.to_timestamp("time")))
# Create the status column
df = df.withColumn("status",
                   F.concat(F.col('world'),
                            F.lit('-'),
                            F.col('level')))
# Apply the window
df = df.withColumn("endtime",
                   F.lead("time").over(w))
# Substract one day from the endtime
df = df.withColumn("endtime",
                   F.expr("date_sub(to_date(endtime), 1)"))
# Fill in the empty `endtime` column with today's date
df = df.withColumn("endtime",
                   F.when(F.col("endtime").isNull(),
                          F.expr("current_date()"))\
                           .otherwise(F.col("endtime")))
# Select the columns and explore the sequence
df.select("character", "status", "time", "endtime").show()

df.createOrReplaceTempView("df")

sc.sql("SELECT status, time, endtime FROM df WHERE character = 'Peach'").show()

sc.sql("SELECT status FROM df WHERE character = 'Peach' AND '2020-05-24' BETWEEN time AND endtime").show()

# Create the events view
eventsdf = sc.createDataFrame(events)
eventsdf.createOrReplaceTempView("events")

sc.sql("""
WITH statuses AS (
    SELECT
        character,
        CONCAT(world, '-', level) AS status,
        to_date(time) AS start,
        to_date(DATE_SUB(LEAD(time) OVER (PARTITION BY character ORDER BY time), 1)) AS end
    FROM
        events
)
SELECT
  character,
  status,
  start,
  IF(end IS NOT NULL, end, current_date()) AS end
FROM statuses
""").show()