SparkML GBTClassifier fails to convert to ONNX

[yungcero]

Hi all - I am running into issues where I am unable to convert a Spark ML GBTClassifier to ONNX. Not sure if this is a initialization issue or something that has not been added in - so would like some insight into a solution or update so that I can help contribute if this feature is not supported yet. The code that I am using below to replicate the error is pulled from the onnxmltools/test folder for spark ml GBT Classifier. I have also replicated this error with other models I have trained where my label columns are double type (0.0/1.0) and my feature columns are created using the VectorAssembler.

lib versions:
onnxmltools: 1.11.2
onnxconverter-common: 1.13.0
pyspark: 3.3.2

from pyspark.ml.linalg import Vectors
raw_data = spark.createDataFrame(
        [(1.0, Vectors.dense(1.0)), (0.0, Vectors.sparse(1, [], []))],
        ["label", "features"],
    )
string_indexer = StringIndexer(inputCol="label", outputCol="indexed")
si_model = string_indexer.fit(raw_data)
data = si_model.transform(raw_data)
gbt = GBTClassifier(maxIter=5, maxDepth=2, labelCol="indexed", seed=42)
model = gbt.fit(data)
feature_count = data.first()[1].size
model_onnx = convert_sparkml(
    model,
    "Sparkml GBT Classifier",
    [("features", FloatTensorType([1, feature_count]))],
    spark_session= spark,
    target_opset=9
)

This will fail on convert_sparkml method, specifically in this class https://github.com/onnx/onnxmltools/blob/main/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py on line 65 when trying to read the written model as a parquet. This is the error message: AnalysisException: Unable to infer schema for Parquet. It must be specified manually.

Any help is appreciated. Thanks in advance!

[yungcero]

Found some bugs when walking through the source code in regards to how spark handles the model conversion - will work a fix and push up a MR for review

[yungcero]

Going to clean up the spark code on our forked repo then will open a pull request to merge back into the main onnxtools

[xadupre]

I'll close this since the PR fixing it was merged.

[yungcero]

Awesome. Thank you and sounds good! @xadupre