TaskNotSerializable with Spark Dataset transformations - scala

Question

I'm trying to filter the Spark dataset by one of the fields but as the result of transformations, I receive TaskNotSerializable. Here's how the method looks like:
def filterData(input: Dataset[CustomType], idsList: List[Int]): Dataset[CustomType] =
input.flatMap { record =>
val filtered = record.data.filter(rec => idsList.contains(rec.id))
if (filtered.nonEmpty)
Seq(record.withFields(filtered))
else
Iterable.empty
}
I tried to do similar transformations on Spark dataframe and it worked without any issues:
input.withColumn("arr", explode($”data”))
.filter($"arr.id".isin(idsList: _*)) //List(34,81,95)
.drop("arr")
.as[CustomType]
How dataset transformations could be fixed to avoid this error?
Custom type structure is the following:
|-- url: string (nullable = true)
|-- data: array (nullable = true)
| |-- element: struct (containsNull = true)
| | |-- id: integer (nullable = false)
| | |-- expiration: long (nullable = false)
| | |-- weight: integer (nullable = false)
|-- queryParams: array (nullable = true)
|-- element: string (containsNull = true)
Error serialization stack looks like this:
Serialization stack:
- object not serializable (class: job.DataSink, value: job.DataSink#6dda8f39)
- field (class: job.DataSink$$anonfun$filterData$1, name: $outer, type: class job.DataSink)
- object (class job.DataSink$$anonfun$filterData$1, <function1>)
- field (class: org.apache.spark.sql.Dataset$$anonfun$flatMap$1, name: func$6, type: interface scala.Function1)
- object (class org.apache.spark.sql.Dataset$$anonfun$flatMap$1, <function1>)
- field (class: org.apache.spark.sql.execution.MapPartitionsExec, name: func, type: interface scala.Function1)
- object (class org.apache.spark.sql.execution.MapPartitionsExec, MapPartitions <function1>, obj#392: schemas.data.CustomType
+- Scan[obj#383]
)

Answer 1

From serialization stack I would say that this function is part of job.DataSink class and Spark is trying to serialize whole class to send it to worker. You need to implement Serializable if you want to make it serializable by Spark.