Have you tried using the pandas udf anotation in Spark instead? I’d guess that could make this obsolete (at least for those models that can be implemented as such), and would wonder about performance comparison