TWS Scala Top K

%sh
pip install dbldatagen

%python
import dbldatagen as dg
import time
from pyspark.sql.types import IntegerType, FloatType, StringType, TimestampType
from datetime import datetime, timedelta

# Set table name
table_name = f"synthetic_data_{int(time.time())}"
dbutils.widgets.text("table_name", table_name)
print(table_name)

# Generate session data with user_id, action_type, and timestamp
data_rows = 100
now = datetime.now()
one_hour_ago = now - timedelta(hours=1)

df_spec = (dg.DataGenerator(spark, name="session_data", rows=data_rows)
           .withColumn("user", StringType(), values=['user1', 'user2', 'user3', 'user4', 'user5'], random=True)
           .withColumn("value", IntegerType(), minValue=0, maxValue=200)
)

                            
df = df_spec.build()

# Write to Delta table
df.write.format("delta").mode("overwrite").saveAsTable(table_name)

// Scala code
spark.conf.set(
  "spark.sql.streaming.stateStore.providerClass",
  "com.databricks.sql.streaming.state.RocksDBStateStoreProvider"
)

val tableName = dbutils.widgets.get("table_name")
// Use spark.table() instead of read.format("delta").load()
val df = spark.table(tableName)
display(df)

// Import the RocksDB state store provider
import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
import java.util.UUID
import org.apache.spark.sql.streaming.StatefulProcessor
import org.apache.spark.sql.streaming._

import java.sql.Timestamp
import org.apache.spark.sql.Encoders

import scala.collection.mutable.PriorityQueue

class TopKStatefulProcessor(k: Int)
  extends StatefulProcessor[String, (String, Int), (Int, String, Int)] {
  // _topKState remembers the top-k elements from the previous micro-batch. It may have
  // less than k elements, but it will never have more than k elements.
  @transient private var _topKState: ListState[(String, Int)] = _

  override def init(
      outputMode: OutputMode,
      timeMode: TimeMode): Unit = {
    _topKState = getHandle.getListState[(String, Int)](
      "topKState", Encoders.tuple(Encoders.STRING, Encoders.scalaInt), TTLConfig.NONE)
  }
  override def handleInputRows(
      key: String,
      inputRows: Iterator[(String, Int)],
      timerValues: TimerValues)
      : Iterator[(Int, String, Int)] = {
    // Keep a heap of the top k elements, ordered by the Int field of (String, Int).
    val pq = new PriorityQueue[(String, Int)]()(Ordering.by(-_._2))

    inputRows.foreach { case (name, value) =>
      if (pq.size < k) {
        pq.enqueue((name, value))
      } else {
        // If the new element is larger than the minimum, remove the
        // minimum and add the new element.
        if (value > pq.head._2) {
          println(s"Adding to pq: ($name, $value) since $value > ${pq.head._2}")
          pq.dequeue()
          pq.enqueue((name, value))
        } else {
          println(s"Not adding to pq: ($name, $value) since $value <= ${pq.head._2}")
        }
      }
    }

    // Merge the top-k elements from the previous micro-batch with the top-k elements
    // from the current micro-batch.
    // For loop of k elements
    val immutableOldTopKState = _topKState.get().toList

    var oldTopKState = immutableOldTopKState
    // Ordered from largest to smallest
    var newTopKState = Array[(String, Int)]()

    var heapEntries = pq.toArray.sortBy(-_._2)

    var bothEmpty = false
    while (newTopKState.size < k && !bothEmpty) {
      // Grab the first of each side
      val topHeapElement = heapEntries.headOption
      val topStateElement = oldTopKState.headOption

      // If both are empty, we're done
      (topHeapElement, topStateElement) match {
        case (None, None) => bothEmpty = true
        case (Some(heapElement), None) =>
          newTopKState = newTopKState :+ heapElement
          heapEntries = heapEntries.tail
        case (None, Some(stateElement)) =>
          newTopKState = newTopKState :+ stateElement
          oldTopKState = oldTopKState.tail
        case (Some(heapElement), Some(stateElement)) =>
          // Prioritize the heap element if it is the same value as the state element
          if (heapElement._2 >= stateElement._2) {
            newTopKState = newTopKState :+ heapElement
            heapEntries = heapEntries.tail
          } else {
            newTopKState = newTopKState :+ stateElement
            oldTopKState = oldTopKState.tail
          }
      }
    }
    _topKState.put(newTopKState)

    // Emit only the top-k elements that changed.
    newTopKState.zipWithIndex.flatMap { case (record, rank) =>
      if (immutableOldTopKState.size <= rank || immutableOldTopKState(rank) != record) {
        Seq((rank + 1, record._1, record._2))
      } else {
        Seq.empty
      }
    }.iterator
  }
}

val inputStream = spark.readStream
  .format("delta")
  .option("maxFilesPerTrigger", "1")
  .table(tableName)  // Use the table name we created
  .as[(String, Int)]

val baseLocation = "/Workspace/Users/bo.gao@databricks.com/tws/" + UUID.randomUUID().toString
val checkpointLocation = baseLocation + "/checkpoint"
val outputTable = baseLocation + "/output"

import spark.implicits._
import org.apache.spark.sql.streaming.Trigger

val result = inputStream
  .groupByKey(x => x._1)
  .transformWithState[(Int, String, Int)](new TopKStatefulProcessor(5),
     TimeMode.None(),
     OutputMode.Update())

val query = result.writeStream
  .format("console")
  .trigger(Trigger.AvailableNow())
  .option("checkpointLocation", checkpointLocation)
  .outputMode("update")
  .start()