[Python] Retrieve categories from Spark ML OneHotEncoder?

I've used a OneHotEncoder in a Spark ML pipeline:

from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler

schema = StructType(
[StructField("PassengerId", DoubleType()),
StructField("Survived", DoubleType()),
StructField("Pclass", DoubleType()),
StructField("Name", StringType()),
StructField("Sex", StringType()),
StructField("Age", DoubleType()),
StructField("SibSp", DoubleType()),
StructField("Parch", DoubleType()),
StructField("Ticket", StringType()),
StructField("Fare", DoubleType()),
StructField("Cabin", StringType()),
StructField("Embarked", StringType())
])

titanic = (
spark
.read
.option("header", "true")
.schema(schema)
.csv("/working_dir/data/titanic.csv")
.na.fill(0)
.na.fill("Nulo")
)
trainDF, testDF = titanic.randomSplit([0.8, 0.2], seed=42)

indexers = [StringIndexer(inputCol=c, outputCol=c + "_index", handleInvalid="keep") for c in ["Sex", "Embarked"]]
ohe = OneHotEncoder(inputCols=[indexer.getOutputCol() for indexer in indexers], outputCols=[f"{indexer.getInputCol()}_onehot" for indexer in indexers])
vectorAssembler = (
VectorAssembler()
.setInputCols(["Sex_onehot", "Embarked_onehot", "Pclass", "Age", "Fare"])
.setOutputCol("features")
)
from pyspark.ml import Pipeline
pipeline = Pipeline(stages=[
*indexers,
ohe,
vectorAssembler
])
trainDF, testDF = titanic.randomSplit([0.8, 0.2], seed=42)
fitted_pipeline = pipeline.fit(trainDF)


Now, I would like to retrieve the categories created for each variable, in a similar fashion as with sklearn:

>>> enc.get_feature_names_out(['gender', 'group'])
array(['gender_Female', 'gender_Male', 'group_1', 'group_2', 'group_3'], ...)


I've tried printing the ParamMap to see if I could extract them from there, with no luck. Is there any way to obtain something similar to what I'm looking for?

Continue reading...