Feature Selection in PySpark

PySpark

Feature Selection

Machine Learning

Data Processing

Big Data

Feature Selection in PySpark

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Machine Learning Big Data

Introduction

Feature selection in PySpark reduces the number of input variables to improve model performance and reduce training time on large datasets. PySpark's MLlib provides VectorSlicer for index-based selection, ChiSqSelector for statistical selection, and UnivariateFeatureSelector (Spark 3.1+) for flexible statistical tests. For correlation-based filtering and variance thresholds, you combine PySpark DataFrame operations with these selectors.

VectorSlicer (Index-Based Selection)

VectorSlicer selects features by their index positions within a feature vector:

python

1from pyspark.ml.feature import VectorSlicer, VectorAssembler
2from pyspark.sql import SparkSession
3
4spark = SparkSession.builder.appName("FeatureSelection").getOrCreate()
5
6df = spark.createDataFrame([
7    (1.0, 2.0, 3.0, 4.0, 0),
8    (5.0, 6.0, 7.0, 8.0, 1),
9    (9.0, 10.0, 11.0, 12.0, 0),
10], ["f1", "f2", "f3", "f4", "label"])
11
12# Assemble features into a single vector column
13assembler = VectorAssembler(inputCols=["f1", "f2", "f3", "f4"], outputCol="features")
14df = assembler.transform(df)
15
16# Select features at index 0 and 2 (f1 and f3)
17slicer = VectorSlicer(inputCol="features", outputCol="selected_features", indices=[0, 2])
18result = slicer.transform(df)
19result.select("selected_features", "label").show()
20# +------------------+-----+
21# | selected_features|label|
22# +------------------+-----+
23# |       [1.0, 3.0] |    0|
24# |       [5.0, 7.0] |    1|
25# |       [9.0, 11.0]|    0|
26# +------------------+-----+

Use VectorSlicer when you already know which feature indices to keep (e.g., from prior analysis).

ChiSqSelector (Statistical Selection)

ChiSqSelector uses the chi-squared test to select features most correlated with the label (for classification tasks):

python

1from pyspark.ml.feature import ChiSqSelector, VectorAssembler
2
3df = spark.createDataFrame([
4    (1.0, 0.0, 3.0, 1.0, 1),
5    (2.0, 1.0, 3.0, 2.0, 0),
6    (3.0, 0.0, 3.0, 3.0, 1),
7    (4.0, 1.0, 3.0, 4.0, 0),
8    (5.0, 0.0, 3.0, 5.0, 1),
9], ["f1", "f2", "f3", "f4", "label"])
10
11assembler = VectorAssembler(inputCols=["f1", "f2", "f3", "f4"], outputCol="features")
12df = assembler.transform(df)
13
14# Select top 2 features by chi-squared statistic
15selector = ChiSqSelector(
16    numTopFeatures=2,
17    featuresCol="features",
18    outputCol="selected_features",
19    labelCol="label"
20)
21
22model = selector.fit(df)
23result = model.transform(df)
24result.select("selected_features", "label").show()
25
26# See which feature indices were selected
27print("Selected feature indices:", model.selectedFeatures)

Selection Types

python

1# Select by number of top features
2ChiSqSelector(numTopFeatures=5, selectorType="numTopFeatures")
3
4# Select by percentile (top 50% of features)
5ChiSqSelector(percentile=0.5, selectorType="percentile")
6
7# Select by false positive rate threshold
8ChiSqSelector(fpr=0.05, selectorType="fpr")

UnivariateFeatureSelector (Spark 3.1+)

UnivariateFeatureSelector supports multiple statistical tests depending on feature and label types:

python

1from pyspark.ml.feature import UnivariateFeatureSelector
2
3selector = UnivariateFeatureSelector(
4    featuresCol="features",
5    outputCol="selected_features",
6    labelCol="label",
7    selectionMode="numTopFeatures"
8)
9
10selector.setFeatureType("continuous")
11selector.setLabelType("categorical")
12selector.setSelectionThreshold(3)  # top 3 features
13
14model = selector.fit(df)
15result = model.transform(df)

Feature Type	Label Type	Test Used
continuous	categorical	ANOVA F-test
categorical	categorical	Chi-squared
continuous	continuous	F-regression

Correlation-Based Filtering

Remove highly correlated features to reduce redundancy:

python

1from pyspark.ml.stat import Correlation
2from pyspark.ml.feature import VectorAssembler
3import numpy as np
4
5assembler = VectorAssembler(inputCols=["f1", "f2", "f3", "f4"], outputCol="features")
6df_vec = assembler.transform(df)
7
8# Compute Pearson correlation matrix
9corr_matrix = Correlation.corr(df_vec, "features", "pearson").head()[0]
10corr_array = corr_matrix.toArray()
11
12# Find highly correlated feature pairs (threshold > 0.9)
13n_features = corr_array.shape[0]
14features_to_drop = set()
15
16for i in range(n_features):
17    for j in range(i + 1, n_features):
18        if abs(corr_array[i][j]) > 0.9:
19            features_to_drop.add(j)  # Drop the later feature
20
21print(f"Features to drop (indices): {features_to_drop}")
22
23# Use VectorSlicer to keep remaining features
24keep_indices = [i for i in range(n_features) if i not in features_to_drop]
25slicer = VectorSlicer(inputCol="features", outputCol="filtered_features", indices=keep_indices)

Variance Threshold

Remove features with near-zero variance (they carry no information):

python

1from pyspark.sql import functions as F
2
3feature_cols = ["f1", "f2", "f3", "f4"]
4threshold = 0.01
5
6# Calculate variance for each feature
7variances = {}
8for col_name in feature_cols:
9    var_value = df.select(F.variance(col_name)).collect()[0][0]
10    variances[col_name] = var_value
11
12# Keep features above threshold
13selected = [col for col, var in variances.items() if var > threshold]
14print(f"Selected features: {selected}")
15
16# Reassemble with selected features
17assembler = VectorAssembler(inputCols=selected, outputCol="features")
18df_filtered = assembler.transform(df)

Feature Selection in a Pipeline

python

1from pyspark.ml import Pipeline
2from pyspark.ml.feature import VectorAssembler, ChiSqSelector, StringIndexer
3from pyspark.ml.classification import RandomForestClassifier
4
5pipeline = Pipeline(stages=[
6    StringIndexer(inputCol="category", outputCol="label"),
7    VectorAssembler(inputCols=["f1", "f2", "f3", "f4"], outputCol="raw_features"),
8    ChiSqSelector(numTopFeatures=2, featuresCol="raw_features",
9                  outputCol="features", labelCol="label"),
10    RandomForestClassifier(featuresCol="features", labelCol="label")
11])
12
13model = pipeline.fit(train_df)
14predictions = model.transform(test_df)

Common Pitfalls

ChiSqSelector on continuous labels: The chi-squared test is designed for categorical labels. Using it with continuous targets gives meaningless results. Use UnivariateFeatureSelector with featureType="continuous" and labelType="continuous" for regression tasks.
Feature indices change after VectorSlicer: After slicing, feature indices reset to 0, 1, 2, ... in the output vector. If you need to map back to original feature names, maintain a mapping from the assembler's inputCols.
Not assembling features first: PySpark MLlib selectors work on vector columns, not individual DataFrame columns. You must use VectorAssembler to combine individual columns into a feature vector before applying any selector.
Correlation matrix on sparse data: Correlation.corr() converts the feature vector to a dense matrix. For datasets with millions of features, this causes out-of-memory errors. Use sampling or batch processing for very wide datasets.
Applying feature selection after train-test split: Fit the selector on training data only and use model.transform() on test data. Fitting on the full dataset leaks information from the test set into the feature selection process.

Summary

Use VectorSlicer for manual index-based feature selection
Use ChiSqSelector for statistical feature selection in classification tasks
Use UnivariateFeatureSelector (Spark 3.1+) for flexible statistical tests across feature and label types
Compute the correlation matrix with Correlation.corr() and drop redundant features above a threshold
Always assemble individual columns into a vector with VectorAssembler before applying selectors
Integrate feature selection into a Pipeline for reproducible workflows

System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Feature Selection in PySpark

ML System Design practice on Codemia

Introduction

VectorSlicer (Index-Based Selection)

ChiSqSelector (Statistical Selection)

Selection Types

UnivariateFeatureSelector (Spark 3.1+)

Correlation-Based Filtering

Variance Threshold

Feature Selection in a Pipeline

Common Pitfalls

Summary

Related reading

System Design Fundamentals

Track what you have practised

ML System Design practice on Codemia