MLeap Serializing the XGBoost model using PySpark API

Hi,

I was able to serialize rhe XGBoost model with MLeap using the older PySPark API (https://github.com/dmlc/xgboost/pull/4656) as shown below:

import mleap.pyspark
from mleap.pyspark.spark_support import SimpleSparkSerializer

trans_model = model.transform(df)
local_path = "jar:file:/tmp/pyspark.model.zip"
model.serializeToBundle(local_path, trans_model)

But I am not able to do with official PySpark API (https://xgboost.readthedocs.io/en/stable/tutorials/spark_estimator.html) support. Is there anything flags that I need to use to make this work?

@hcho3 Any insights on this topic?

@WeichenXu123 Do you have any insights on this issue?