When you first dive into Spark MLlib, the biggest shock isn't usually the math—it's the API. If you're used to standard Scala collections or even basic Spark DataFrames, you'll find that MLlib is very opinionated about how data is structured. It doesn't want a bunch of separate columns for your features; it wants one single column containing a vector of those features.

Creating some synthetic churn data

Let's build a simple model to predict customer churn. I'll start by creating a small dataset. In a real project, you'd be loading this from S3 or a database, but for this example, we'll just hardcode some values so you can run this immediately.

import org.apache.spark.sql.SparkSession
import org.apache.spark.ml.classification.LogisticRegression
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator

val spark = SparkSession.builder().appName("ChurnPrediction").master("local[*]").getOrCreate()
import spark.implicits._

val data = Seq(
  (1, 12.0, 50.0, 0), // Tenure, MonthlyCharge, Churned
  (2, 1.0, 80.0, 1),
  (3, 24.0, 40.0, 0),
  (4, 2.0, 90.0, 1),
  (5, 36.0, 30.0, 0),
  (6, 5.0, 70.0, 1)
).toDF("id", "tenure", "monthly_charge", "label")

The "Vector" hurdle

Here is where I usually trip up when I'm rushing. My instinct is to just throw this DataFrame into the model and tell it which columns are the features. I'll show you what I mean—I tried to do this the first time I wrote this snippet:

// THIS WILL FAIL
val lr = new LogisticRegression()
val model = lr.fit(data) 

If you run that, Spark will throw an AnalysisException complaining that it can't find a column named "features". MLlib algorithms are designed to be generic, so they don't look for "tenure" or "monthly_charge"; they look for exactly one column called `features` (by default) that holds a Vector. To fix this, we have to use the VectorAssembler. Think of it as a "bundler" that squashes multiple numeric columns into a single vector column.

val assembler = new VectorAssembler()
  .setInputCols(Array("tenure", "monthly_charge"))
  .setOutputCol("features")

val vectorizedData = assembler.transform(data)
vectorizedData.show()

Now our DataFrame has a new column where the values look like [12.0, 50.0]. This is exactly what the machine learning models expect.

Training the Logistic Regression model

Now that the data is in the right shape, we can actually train the model. I'm choosing Logistic Regression because it's the bread and butter for binary classification (churn vs. no-churn). I'll explicitly set the label column since we named ours "label" (which happens to be the default, but it's good practice to be explicit).

val lr = new LogisticRegression()
  .setLabelCol("label")
  .setFeaturesCol("features")

val lrModel = lr.fit(vectorizedData)

// Let's see how the model performs on the same data
val predictions = lrModel.transform(vectorizedData)
predictions.select("id", "label", "prediction").show()

Evaluating the results

Looking at the table is fine for five rows, but for anything real, you need a metric. I'll use the BinaryClassificationEvaluator. By default, this calculates the Area Under the ROC curve (AUC). An AUC of 1.0 is perfect; 0.5 is basically guessing.

val evaluator = new BinaryClassificationEvaluator()
  .setLabelCol("label")
  .setRawPredictionCol("probability")

val auc = evaluator.evaluate(predictions)
println(s"Model AUC: $auc")

It's a simple pipeline: Data → VectorAssembler → Model → Evaluator. Once you memorize that sequence, using MLlib from Scala becomes much less intimidating.

📋 Practical Task

Build a House Price Linear Regression Pipeline

In this exercise, you will move from classification to regression. Your goal is to predict a continuous house price based on two features: sqft and num_rooms.

• Create a synthetic DataFrame with at least 5 rows containing sqft (Double), num_rooms (Double), and price (Double).
• Use a VectorAssembler to combine sqft and num_rooms into a single features column.
• Use org.apache.spark.ml.regression.LinearRegression to train a model using price as the label.
• Transform the data to get predictions and print the final DataFrame showing the actual price vs. the predicted price.