Scala spark aggregator NotSerializableException - scala

Question

I have DataSet with id and sub_ids. I want to use Aggregator for grouping sub_ids per ids into set.
When I execute below code in databricks:
%scala
import org.apache.spark.sql.{DataFrame, Dataset, Encoder}
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
case class SubIdPerId(id: Int, sub_id: String)
case class SubIdsPerId(id: Int, sub_ids: Set[String])
object PreprocessData{
def createSubIdsArray(df: Dataset[SubIdPerId]): Dataset[SubIdsPerId] = {
val transformedData: Dataset[SubIdsPerId] =
df.groupByKey(_.id)
.agg(
PreprocessData.distinctSubIdsAggregator.toColumn.name("sub_ids")
).map{
case(id: Int, si: Set[String]) => SubIdsPerId(id, si)
}
transformedData
}
val distinctSubIdsAggregator: Aggregator[SubIdPerId, Set[String], Set[String]] = new Aggregator[SubIdPerId, Set[String], Set[String]] {
override def zero: Set[String] = Set[String]()
override def reduce(es: Set[String], supi: SubIdPerId): Set[String] =
es + supi.sub_id
override def merge(wx: Set[String], wy: Set[String]): Set[String] =
wx.union(wy)
override def finish(reduction: Set[String]): Set[String] = reduction
override def bufferEncoder: Encoder[Set[String]] =
implicitly(ExpressionEncoder[Set[String]])
override def outputEncoder: Encoder[Set[String]] =
implicitly(ExpressionEncoder[Set[String]])
}
}
val df_t = Seq(
(1, "zk67"),
(1, "gg89"),
(2, "gg97"),
(2, "gd01"),
(2, "af83"),
(3, "af84"),
(3, "gd77"),
(3, "gd73"),
(3, "cl55"),
(3, "zk67")
).toDF("id","sub_id")
val ds_t = df_t.as[SubIdPerId]
val transformedData = PreprocessData.createSubIdsArray(ds_t)
display(transformedData)
I receive:
NotSerializableException: $line813d300e73a74f3e9ab1e2bfe2afe2de59.$read$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw$PreprocessData$
Additional note:
When I execute this code without wrapping it into object it works:
ds_t.groupByKey(_.id).agg(distinctSubIdsAggregator.toColumn.name("sub_ids"))