|
| 1 | +from pyspark.sql import SparkSession |
| 2 | +from pyspark.sql.functions import from_json, col, to_timestamp, window |
| 3 | +from pyspark.sql.types import StructType, StructField, StringType, DoubleType |
| 4 | + |
| 5 | +from lib.logger import Log4j |
| 6 | + |
| 7 | +if __name__ == "__main__": |
| 8 | + spark = SparkSession \ |
| 9 | + .builder \ |
| 10 | + .appName("Sliding Window Demo") \ |
| 11 | + .master("local[3]") \ |
| 12 | + .config("spark.streaming.stopGracefullyOnShutdown", "true") \ |
| 13 | + .getOrCreate() |
| 14 | + |
| 15 | + logger = Log4j(spark) |
| 16 | + |
| 17 | + invoice_schema = StructType([ |
| 18 | + StructField("InvoiceNumber", StringType()), |
| 19 | + StructField("CreatedTime", StringType()), |
| 20 | + StructField("StoreID", StringType()), |
| 21 | + StructField("TotalAmount", DoubleType()) |
| 22 | + ]) |
| 23 | + |
| 24 | + kafka_df = spark.readStream \ |
| 25 | + .format("kafka") \ |
| 26 | + .option("kafka.bootstrap.servers", "localhost:9092") \ |
| 27 | + .option("subscribe", "invoices") \ |
| 28 | + .option("startingOffsets", "earliest") \ |
| 29 | + .load() |
| 30 | + |
| 31 | + value_df = kafka_df.select(from_json(col("value").cast("string"), invoice_schema).alias("value")) |
| 32 | + |
| 33 | + # value_df.printSchema() |
| 34 | + # value_df.show(truncate=False) |
| 35 | + |
| 36 | + invoice_df = value_df.select("value.*") \ |
| 37 | + .withColumn("CreatedTime", to_timestamp("CreatedTime", "yyyy-MM-dd HH:mm:ss")) |
| 38 | + |
| 39 | + count_df = invoice_df.groupBy("StoreID", |
| 40 | + window("CreatedTime", "5 minute", "1 minute")).count() |
| 41 | + |
| 42 | + # count_df.printSchema() |
| 43 | + # count_df.show(truncate=False) |
| 44 | + |
| 45 | + output_df = count_df.select("StoreID", "window.start", "window.end", "count") |
| 46 | + |
| 47 | + windowQuery = output_df.writeStream \ |
| 48 | + .format("console") \ |
| 49 | + .outputMode("update") \ |
| 50 | + .option("checkpointLocation", "chk-point-dir") \ |
| 51 | + .trigger(processingTime="1 minute") \ |
| 52 | + .start() |
| 53 | + |
| 54 | + logger.info("Counting Invoices") |
| 55 | + windowQuery.awaitTermination() |
0 commit comments