PySpark Structured Streaming

from pyspark.sql import SparkSession, Row
import pyspark.sql
import json
from pyspark.sql.types import *
import time
from IPython.display import display, clear_output
import pandas as pd

spark = SparkSession \
    .builder \
    .appName("StructuredStreaming") \
    .getOrCreate()

path = "data/"

# Create DataFrame representing data in the JSON files
def loadData(path):
    #path = "testdata/"
    jsonSchema = (StructType()
                 .add("time",TimestampType())
                 .add("animal_name",StringType())
                 .add("weather",StringType())
                 .add("battery",DoubleType())
                 )
    outdf = spark.read.json(path, schema = jsonSchema).dropna("any")
   
    return outdf

#example print
loadData(path).show()

'''loadData tests'''

cols = StructType([ StructField("time", TimestampType(), True),
                    StructField("animal_name", StringType(), True),
                    StructField("weather", StringType(), True),
                    StructField("battery", DoubleType(), True)])


from datetime import datetime
testTs = datetime(2020, 1, 1)

fakeData = [(testTs, "dog", "cloudy", 100.0)]

fakeDf = spark.createDataFrame(fakeData, cols)

df = loadData(path)

assert df.dtypes == fakeDf.dtypes, "the schema was expected to be %s but it was %s" % (fakeDf.dtypes, df.dtypes)

def animalCount(df):
    df.createOrReplaceTempView("ntable")
    newdf = spark.sql("""
        SELECT animal_name, COUNT(animal_name) AS count FROM ntable
        GROUP BY animal_name
        ORDER BY count DESC
    """)
    return newdf
    # YOUR CODE HERE
    #raise NotImplementedError()

#example print
animalCount(loadData(path)).show()

'''animalCount tests'''

cols = StructType([ StructField("animal_name", StringType(), True),
                    StructField("count", LongType(), False)])


fakeData = [("dog", 1)]

fakeDf = spark.createDataFrame(fakeData, cols)

df = animalCount(loadData(path))

assert df.dtypes == fakeDf.dtypes, "the schema was expected to be %s but it was %s" % (fakeDf.dtypes, df.dtypes)

assert df.count() == 5, "the number of rows was expected to be 5 but it was %s" % df.count()

df = df.toPandas()

assert df.loc[0][1] >= df.loc[1][1], "the first item was expected to have higher count than the second"
assert df.loc[3][1] >= df.loc[4][1], "the fourth item was expected to have higher count than the last"
assert df.loc[0][0] == "Rabbit", "the first item was expected to be Rabbit but it was %s" % df.loc[0][0]
assert df.loc[4][0] == "Wolf", "the last item was expected to be Wolf but it was %s" % df.loc[4][0]

def inputDf(path):  
    jsonSchema = (StructType()
                 .add("time",TimestampType())
                 .add("animal_name",StringType())
                 .add("weather",StringType())
                 .add("battery",DoubleType())
                 )
    inputDF = (spark
               .readStream 
               .schema(jsonSchema)
               .option("maxFilesPerTrigger",1)
               .json(path)
               .dropna("any")
               
              )
    # YOUR CODE HERE
    return inputDF

def outputDf(inputDF):
    filterDF = inputDF.createOrReplaceTempView("streamTable")
    newdf = spark.sql("""
        SELECT animal_name, COUNT(animal_name) AS count 
        FROM streamTable
        GROUP BY animal_name
        ORDER BY count DESC
    """)
    return newdf
    # YOUR CODE HERE
    #raise NotImplementedError()

def createQuery(outputDF):
    query = (outputDF
             .writeStream 
             .format("memory")
             .outputMode("complete")
             .queryName("counts")
             .start()
            )
    # YOUR CODE HERE
    time.sleep(0.5)
    return query

'''streaming tests'''
inputStreamDf = inputDf(path)
outputStreamDf= outputDf(inputStreamDf)
query = createQuery(outputStreamDf)

assert outputStreamDf.isStreaming, "the outputDF was expected to be streaming"

df = spark.sql("select * from counts")

assert df.dtypes == fakeDf.dtypes, "the schema was expected to be %s but it was %s" % (fakeDf.dtypes, df.dtypes)

status = {'message': 'Processing new data',
 'isDataAvailable': True,
 'isTriggerActive': True}

assert query.status == status, "the status was expected to be %s but it was %s" % (status, query.status)

assert df.count() == 0, "the number of rows was expected to be 0 when the streaming just started but it was %s" % df.count()
time.sleep(20)
assert df.count() == 5, "the number of rows was expected to be 5 when the streaming has been going for a while but it was %s" % df.count()

df = df.toPandas()

assert df.loc[0][1] >= df.loc[1][1], "the first item was expected to have higher count than the second"
assert df.loc[3][1] >= df.loc[4][1], "the fourth item was expected to have higher count than the last"
assert df.loc[0][0] == "Rabbit" or df.loc[0][0] == "Squirrel", "the first item was expected to be Rabbit or Squirrel but it was %s" % df.loc[0][0]
assert df.loc[4][0] == "Wolf" or df.loc[4][0] == "Bear", "the last item was expected to be Wolf or Bear but it was %s" % df.loc[4][0]

# stream printing in real-time
# you can increase the total streaming simulation duration by increasing "nIter"
# However, when finally submitting the assignment make sure to change "nIter" back to a smaller value. 
nIter = 5
for i in range(nIter):
    clear_output(wait=True)
    display(query.status)
    display()
    time.sleep(3)

spark.stop()