PySpark - OneHotEncoder

 

from pyspark.ml.feature import OneHotEncoderdata = [(0.0, 1.0), (1.0, 0.0), (2.0, 1.0)]
columns = ["input1", "input2"]
df = spark.createDataFrame(data, columns)encoder = OneHotEncoder(inputCols=["input1", "input2"],outputCols=["output1", "output2"]
)encoder_model = encoder.fit(df)
encoded_df = encoder_model.transform(df)
encoded_df.show(truncate=False)

 

+------+------+-------------+-------------+
|input1|input2|output1      |output2      |
+------+------+-------------+-------------+
|0.0   |1.0   |(2,[0],[1.0])|(1,[],[])    |
|1.0   |0.0   |(2,[1],[1.0])|(1,[0],[1.0])|
|2.0   |1.0   |(2,[],[])    |(1,[],[])    |
+------+------+-------------+-------------+

 

 

OneHotEncoder output format

Each encoded column is a SparseVector, shown as:

(size, [indices], [values])

Where:

• size = number of categories minus 1
  (because OneHotEncoder drops the last category to avoid multicollinearity)

• indices = which category this row belongs to

• values = always [1.0]

---

📌 Vocabulary sizes in your case

From the output:

For output1

(2,[0],[1.0]) → size = 2  

Meaning input1 has 3 categories:

0.0 → index 0
1.0 → index 1
2.0 → index dropped (all zeros)

For output2

(1,[0],[1.0]) → size = 1

Meaning input2 has 2 categories:

0.0 → index 0
1.0 → dropped (all zeros)

 

---

🔍 Now decode your rows

---

▶ Row 1

input1 = 0.0
input2 = 1.0output1 = (2,[0],[1.0])
output2 = (1,[],[])

• input1=0 → category index 0 → one-hot at index 0
  → (2,[0],[1.0])

• input2=1 is the last category, so encoder drops it
  → all zeros (1,[],[])

---

▶ Row 2

input1 = 1.0
input2 = 0.0output1 = (2,[1],[1.0])
output2 = (1,[0],[1.0])

• input1=1 → index 1 → (2,[1],[1.0])

• input2=0 → index 0 → (1,[0],[1.0])

---

▶ Row 3

input1 = 2.0
input2 = 1.0output1 = (2,[],[])
output2 = (1,[],[])

• input1=2 → last category, dropped → (2,[],[])

• input2=1 → last category, dropped → (1,[],[])

---

🎯 Summary

output1

• 3 categories in input1 → output vector size = 2

• Encoded as:

  input1=0.0 → (2,[0],[1.0])
  input1=1.0 → (2,[1],[1.0])
  input1=2.0 → (2,[],[])     # dropped category

output2

• 2 categories in input2 → output vector size = 1

• Encoded as:

  input2=0.0 → (1,[0],[1.0])
  input2=1.0 → (1,[],[])     # dropped category

---

If you want, I can also show you how to:

✔ keep all categories (no drop)
✔ extract the category → vector index mapping
✔ convert SparseVector to dense vector for readability

 

You said:

Yes, I want.

 

ChatGPT said:

Great — here are the three things you asked for, clearly explained with ready-to-use PySpark code:

✔ (A) Keep all categories (no drop)
✔ (B) Get the category → vector index mapping
✔ (C) Convert SparseVector to DenseVector for readability

---

🅰 A. Keep all categories (no drop)

By default, PySpark OneHotEncoder drops the last category to avoid multicollinearity (dummy variable trap).

To keep all categories:

from pyspark.ml.feature import OneHotEncoderencoder = OneHotEncoder(inputCols=["input1", "input2"],outputCols=["output1", "output2"],dropLast=False   # ← keep every category
)encoder_model = encoder.fit(df)
encoded_df = encoder_model.transform(df)
encoded_df.show(truncate=False)

If a column has k categories, the output vector will have k elements (instead of k-1).

---

🅱 B. Get category → vector index mapping

After fitting:

encoder_model.categorySizes

This returns the number of categories for each input column.

Example output:

[3, 2]

Meaning:

• input1 has 3 distinct categories → output1 vector length = 2 (or 3 if dropLast=False)

• input2 has 2 distinct categories → output2 vector length = 1 (or 2 if dropLast=False)

You can also get full mapping:

1. Extract labels for input columns

If input columns were originally string -> indexed -> encoded,
you can inspect the StringIndexer model.

But since your inputs were numeric categories, Spark assumes the category index = the value itself.

So the mapping is simply:

Category value → vector index

Unless categories are not dense (e.g., 10, 20, 50), in which case Spark remaps them internally.

To see mapping explicitly:

for i, size in enumerate(encoder_model.categorySizes):print(f"Column {i} categories: 0..{size - 1}")

Example output:

Column 0 categories: 0..2
Column 1 categories: 0..1

If you want a fully explicit mapping:

for col, size in zip(["input1", "input2"], encoder_model.categorySizes):print(f"Mapping for {col}:")for idx in range(size):print(f"  category {idx} → vector index {idx}")

 

---

🅲 C. Convert SparseVector to DenseVector

Sometimes you want to see the full vector (e.g., [0,1,0]) instead of:

(3,[1],[1.0])

Use a UDF:

from pyspark.ml.linalg import DenseVectorto_dense = F.udf(lambda v: DenseVector(v.toArray()), DenseVector())df_dense = encoded_df.withColumn("output1_dense", to_dense("output1")) \.withColumn("output2_dense", to_dense("output2"))df_dense.show(truncate=False)

 

Example output

 

+------+------+-------------+-------------+-------------+-------------+ |input1|input2|output1 |output2 |output1_dense|output2_dense| +------+------+-------------+-------------+-------------+-------------+ |0.0 |1.0 |(2,[0],[1.0])|(1,[],[]) |[1.0,0.0] |[0.0] | |1.0 |0.0 |(2,[1],[1.0])|(1,[0],[1.0])|[0.0,1.0] |[1.0] | |2.0 |1.0 |(2,[],[]) |(1,[],[]) |[0.0,0.0] |[0.0] | +------+------+-------------+-------------+-------------+-------------+

This makes the encoding much easier to inspect.

---

🎉 Summary

✔ Keep all categories

 

dropLast=False

✔ See number of categories per column

 

encoder_model.categorySizes

✔ Convert SparseVector → DenseVector

 

DenseVector(v.toArray())