How to map features from the output of a VectorAssembler back to the column names in Spark ML?

by

As of today Spark doesn’t provide any method that can do it for you, so if you have to create your own. Let’s say your data looks like this:

import random
random.seed(1)

df = sc.parallelize([(
    random.choice([0.0, 1.0]), 
    random.choice(["a", "b", "c"]),
    random.choice(["foo", "bar"]),
    random.randint(0, 100),
    random.random(),
) for _ in range(100)]).toDF(["label", "x1", "x2", "x3", "x4"])

and is processed using following pipeline:

from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler
from pyspark.ml import Pipeline
from pyspark.ml.regression import LinearRegression

indexers = [
  StringIndexer(inputCol=c, outputCol="{}_idx".format(c)) for c in ["x1", "x2"]]
encoders = [
    OneHotEncoder(
        inputCol=idx.getOutputCol(),
        outputCol="{0}_enc".format(idx.getOutputCol())) for idx in indexers]
assembler = VectorAssembler(
    inputCols=[enc.getOutputCol() for enc in encoders] + ["x3", "x4"],
    outputCol="features")

pipeline = Pipeline(
    stages=indexers + encoders + [assembler, LinearRegression()])
model = pipeline.fit(df)

Get the LinearRegressionModel:

lrm = model.stages[-1]

Transform the data:

transformed =  model.transform(df)

Extract and flatten ML attributes:

from itertools import chain

attrs = sorted(
    (attr["idx"], attr["name"]) for attr in (chain(*transformed
        .schema[lrm.summary.featuresCol]
        .metadata["ml_attr"]["attrs"].values())))

and map to the output:

[(name, lrm.summary.pValues[idx]) for idx, name in attrs]

[('x1_idx_enc_a', 0.26400012641279824),
 ('x1_idx_enc_c', 0.06320192217171572),
 ('x2_idx_enc_foo', 0.40447778902400433),
 ('x3', 0.1081883594783335),
 ('x4', 0.4545851609776568)]

[(name, lrm.coefficients[idx]) for idx, name in attrs]

[('x1_idx_enc_a', 0.13874401585637453),
 ('x1_idx_enc_c', 0.23498565469334595),
 ('x2_idx_enc_foo', -0.083558932128022873),
 ('x3', 0.0030186112903237442),
 ('x4', -0.12951394186593695)]

More Related Contents:

  1. How to split Vector into columns – using PySpark
  2. Tuning parameters for implicit pyspark.ml ALS matrix factorization model through pyspark.ml CrossValidator
  3. How to create a custom Estimator in PySpark
  4. How do I convert an array (i.e. list) column to Vector
  5. Applying UDFs on GroupedData in PySpark (with functioning python example)
  6. creating spark data structure from multiline record
  7. Encode and assemble multiple features in PySpark
  8. How to change a dataframe column from String type to Double type in PySpark?
  9. Retrieve top n in each group of a DataFrame in pyspark
  10. Pyspark: explode json in column to multiple columns
  11. AttributeError: ‘DataFrame’ object has no attribute ‘map’
  12. How to use a Scala class inside Pyspark
  13. How to add third-party Java JAR files for use in PySpark
  14. Spark: Broadcast variables: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transforamtion
  15. Filter Pyspark dataframe column with None value
  16. Cast column containing multiple string date formats to DateTime in Spark
  17. How to explode multiple columns of a dataframe in pyspark
  18. Convert spark DataFrame column to python list
  19. PySpark: multiple conditions in when clause
  20. Count number of non-NaN entries in each column of Spark dataframe in PySpark
  21. Load CSV file with PySpark
  22. Cannot find col function in pyspark
  23. What is the best way to remove accents with Apache Spark dataframes in PySpark?
  24. ImportError: No module named numpy on spark workers
  25. How to pass a constant value to Python UDF?
  26. Add column sum as new column in PySpark dataframe
  27. Pyspark – converting json string to DataFrame
  28. Spark DAG differs with ‘withColumn’ vs ‘select’
  29. How to drop all columns with null values in a PySpark DataFrame?
  30. Random numbers generation in PySpark

Leave a Comment