Spark Catalyst flatMapGroupsWithState: Group State with sorted collection - scala

Question

I am trying to have a sorted collection in the state of my groups and I get an error from catalyst which I think regards default instance creation for the collection.
Below is simplified pipeline that demonstrates the error:
package com.example
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.streaming.{GroupState, GroupStateTimeout, OutputMode, Trigger}
import scala.collection.immutable.TreeMap
case class Event
(
key: String
)
case class KeyState
(
prop: TreeMap[Long, String]
)
object CatalystIssue {
def updateState(k: String, vs: Iterator[Event],
state: GroupState[KeyState]) : Iterator[Event] = vs
def main(args: Array[String]) {
val spark = SparkSession.builder()
.master("local[*]")
.appName("CatalystIssue")
.getOrCreate()
import spark.implicits._
val df = spark.readStream.format("rate")
.load()
.select(lit("a").as("key"))
.as[Event]
.groupByKey(_.key)
.flatMapGroupsWithState(OutputMode.Append(),
GroupStateTimeout.NoTimeout())(updateState)
val query = df.writeStream.format("console")
.trigger(Trigger.ProcessingTime("30 seconds")).start()
query.awaitTermination()
}
}
Which produces the error:
ERROR org.apache.spark.sql.catalyst.expressions.codegen.CodeGenerator - failed to compile: org.codehaus.commons.compiler.CompileException: File 'generated.java', Line 53, Column 106: No applicable constructor/method found for zero actual parameters; candidates are: "public scala.collection.mutable.Builder scala.collection.generic.SortedMapFactory.newBuilder(scala.math.Ordering)"
This might be because Sorted Maps are not supported as a dataframe attribute type although that's not my intention here and I would have thought the KeyState would have been opaque to spark since you don't actually access it like a dataframe attribute.
While not very attractive one option might be to serialize the sorted set into a byte array which is an attribute of the KeyState. i.e.
case class KeyState
(
prop: Array[Byte]
)
If Java Serialization were used would that preserve the internal tree structure of the TreeMap, so that at least that would not have to be be rebuilt? Are there any alternative serialization technologies that would preserve the structure?
It seems useful to be able to keep some sorted collections in the group state, especially as the computation is supposed to be primarily in memory. Is there something about the way spark works that makes this fundamentally unworkable?