Stratified cross-validation in PySpark

Question

I am using the Apache Spark API in python, PySpark (--version 3.0.0), and would ideally like to perform cross-validation of my labelled data in a stratified manner since my data is highly imbalanced! I am currently using the below module.

from pyspark.ml.tuning import CrossValidator

In scikit-learn this is possible by defining a StratifiedKFold and placing it in the cv argument of any of the grid search function. This ensures that every K fold a given estimator is trained on, contains the labelled data in a proportionally representative way.

Is there an equivalent for this in PySpark?

< rant> I would have loved to asked this question from the Spark team or as an improvement suggestion but their GitHub page does not allow for either bug reports or improvement suggestions and their official page requires the accompanying source code suggestions, which is slightly above my skill set!

user_4461048 · Answer

I believe that Spark ML doesn't supports Stratified cross-validation at the moment.

However, you could take a look at spark-stratifier. It's a Spark ML stratified cross validator component that was open sourced by HackerRank a few years ago [1]. You could take a look at that.

It can be installed by running pip install spark-stratifier.

Some example code to get you going:

from pyspark.ml import Pipeline
from pyspark.ml.linalg import Vectors
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.evaluation import BinaryClassificationEvaluator
from pyspark.ml.tuning import ParamGridBuilder
from pyspark.sql import SparkSession

from spark_stratifier import StratifiedCrossValidator

spark = SparkSession.builder.master("local").getOrCreate()

data = spark.createDataFrame([
    (1.0, Vectors.dense([0.0, 1.1, 0.1])),
    (0.0, Vectors.dense([2.0, 1.0, -1.0])),
    (0.0, Vectors.dense([2.0, 1.3, 1.0])),
    (1.0, Vectors.dense([0.0, 1.2, -0.5]))], ["label", "features"])

lr = LogisticRegression(maxIter=10)

pipeline = Pipeline(stages=[lr])

paramGrid = ParamGridBuilder() \
    .addGrid(lr.regParam, [0.1, 0.01]) \
    .build()

scv = StratifiedCrossValidator(
        estimator=pipeline,
        estimatorParamMaps=paramGrid,
        evaluator=BinaryClassificationEvaluator(),
        numFolds=2
      )

model = scv.fit(data)

[1]: https://github.com/interviewstreet/spark-stratifier

Stratified cross-validation in PySpark

1 Answers