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:
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):
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
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:
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:
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):
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
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
