As you’ve opened a new thread for this, I’ll continue the conversation here instead of the old thread were this is semi-unrelated.
Can you please share the full exception? And also how you defined TrainData ?
This is something that usually needs to be configured on Spark, but maybe you are trying to push too much data though at the same time, so I’ll need a bit more information to be able to tell you what is going on.