Exploring streaming

examples/src/main/kotlin/org/jetbrains/kotlinx/spark/examples/Streaming.kt

@@ -20,42 +20,29 @@
 package org.jetbrains.kotlinx.spark.examples
 
 import org.apache.spark.SparkConf
-import org.apache.spark.api.java.JavaSparkContext
+import org.apache.spark.sql.Dataset
+import org.apache.spark.streaming.Duration
 import org.apache.spark.streaming.Durations
 import org.apache.spark.streaming.api.java.JavaStreamingContext
-import org.jetbrains.kotlinx.spark.api.withSpark
-import scala.Tuple2
-import java.io.Serializable
+import org.jetbrains.kotlinx.spark.api.*
 
-data class Row @JvmOverloads constructor(
-    var word: String = "",
-) : Serializable
+data class TestRow(
+    val word: String,
+)
 
-fun main() = withSpark {
-
-    val context = JavaStreamingContext(
-        SparkConf()
-            .setMaster("local[*]")
-            .setAppName("Test"),
-        Durations.seconds(1),
-    )
-
-    val lines = context.socketTextStream("localhost", 9999)
+fun main() = withSparkStreaming(Durations.seconds(1)) {
 
+    val lines = ssc.socketTextStream("localhost", 9999)
     val words = lines.flatMap { it.split(" ").iterator() }
 
     words.foreachRDD { rdd, time ->
+        val dataframe: Dataset<TestRow> = rdd.map { TestRow(it) }.toDS()
 
-        // todo convert rdd to dataset using kotlin data class?
-
-        val rowRdd = rdd.map { Row(it) }
-
-        val dataframe = spark.createDataFrame(rowRdd, Row::class.java)
-
+        dataframe
+            .groupByKey { it.word }
+            .count()
+            .show()
 
     }
 
-
-    context.start()
-    context.awaitTermination()
 }

kotlin-spark-api/3.2/src/main/kotlin/org/jetbrains/kotlinx/spark/api/SparkHelper.kt

@@ -27,7 +27,10 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.SparkSession.Builder
 import org.apache.spark.sql.UDFRegistration
+import org.apache.spark.streaming.Duration
+import org.apache.spark.streaming.api.java.JavaStreamingContext
 import org.jetbrains.kotlinx.spark.api.SparkLogLevel.ERROR
+import kotlin.math.log
 
 /**
  * Wrapper for spark creation which allows setting different spark params.
@@ -105,17 +108,60 @@ inline fun withSpark(sparkConf: SparkConf, logLevel: SparkLogLevel = ERROR, func
     )
 }
 
+
 /**
- * This wrapper over [SparkSession] which provides several additional methods to create [org.apache.spark.sql.Dataset]
+ * Wrapper for spark streaming creation. `spark: SparkSession` and `ssc: JavaStreamingContext` are provided, started,
+ * awaited, and stopped automatically.
+ *
+ * @param batchDuration The time interval at which streaming data will be divided into batches
+ * @param props spark options, value types are runtime-checked for type-correctness
+ * @param master Sets the Spark master URL to connect to, such as "local" to run locally, "local[4]" to
+ *  run locally with 4 cores, or "spark://master:7077" to run on a Spark standalone cluster. By default, it
+ *  tries to get the system value "spark.master", otherwise it uses "local[*]"
+ * @param appName Sets a name for the application, which will be shown in the Spark web UI.
+ *  If no application name is set, a randomly generated name will be used.
+ * @param logLevel Control our logLevel. This overrides any user-defined log settings.
+ * @param func function which will be executed in context of [KSparkStreamingSession] (it means that `this` inside block will point to [KSparkStreamingSession])
+ * todo: provide alternatives with path instead of batchDuration etc
  */
-class KSparkSession(val spark: SparkSession) {
-
-    val sc: JavaSparkContext = JavaSparkContext(spark.sparkContext)
+@JvmOverloads
+inline fun withSparkStreaming(
+    batchDuration: Duration,
+    props: Map<String, Any> = emptyMap(),
+    master: String = SparkConf().get("spark.master", "local[*]"),
+    appName: String = "Kotlin Spark Sample",
+    logLevel: SparkLogLevel = SparkLogLevel.ERROR,
+    func: KSparkStreamingSession.() -> Unit,
+) {
+    withSpark(
+        props = props,
+        master = master,
+        appName = appName,
+        logLevel = logLevel,
+    ) {
+        val ssc = JavaStreamingContext(sc, batchDuration)
+        KSparkStreamingSession(session = this, ssc = ssc).apply {
+            func()
+            ssc.start()
+            ssc.awaitTermination()
+        }
+    }
+}
 
+/**
+ * This wrapper over [SparkSession] provides several additional methods to create [org.apache.spark.sql.Dataset]
+ */
+open class KSparkSession(val spark: SparkSession, val sc: JavaSparkContext = JavaSparkContext(spark.sparkContext)) {
     inline fun <reified T> List<T>.toDS() = toDS(spark)
     inline fun <reified T> Array<T>.toDS() = spark.dsOf(*this)
     inline fun <reified T> dsOf(vararg arg: T) = spark.dsOf(*arg)
     inline fun <reified T> RDD<T>.toDS() = toDS(spark)
     inline fun <reified T> JavaRDDLike<T, *>.toDS() = toDS(spark)
     val udf: UDFRegistration get() = spark.udf()
 }
+
+/**
+ * This wrapper over [SparkSession] and [JavaStreamingContext] provides several additional methods to create [org.apache.spark.sql.Dataset]
+ */
+class KSparkStreamingSession(session: KSparkSession, val ssc: JavaStreamingContext) : KSparkSession(session.spark, session.sc)
+