Spark DataSet搭建LinearRegression：如何设置特征列X与标签列Y

Hey there! Let's break down exactly how to build your Linear Regression model with your Spark DataSet—this is such a common first step, so you’re already on the right path. Here’s a step-by-step guide with code examples to make it concrete:

First, Quick Prep: Fix Column Names (If Needed) #

Looking at your show() output, it looks like your column names have leading spaces (e.g., D_5 instead of D_5). This can cause errors when referencing columns later, so let’s clean that up first:

Scala Example

// Convert DataSet to DataFrame (Spark ML works best with DataFrames) and clean column names
val cleanedDF = yourDataSet.toDF().columns.foldLeft(yourDataSet.toDF()) { (tempDF, colName) =>
  tempDF.withColumnRenamed(colName, colName.trim)
}

Python Example

# Convert DataSet to DataFrame and trim spaces from all column names
cleaned_df = your_dataset.toDF()
for col_name in cleaned_df.columns:
    cleaned_df = cleaned_df.withColumnRenamed(col_name, col_name.strip())

Step 1: Assemble Your Features into a Single Vector #

Spark ML models require all features to be in a single vector column (not separate columns). We’ll use VectorAssembler to combine all your feature columns (everything except W_12_26) into a column named features.

Step 2: Build & Train the Linear Regression Model #

We’ll wrap the feature assembly and model training into a Pipeline—this makes your workflow reusable and easy to tweak later.

Full Scala Code

import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.evaluation.RegressionEvaluator

// Define which columns are features (exclude the label column W_12_26)
val featureCols = cleanedDF.columns.filter(_ != "W_12_26")

// Create the feature assembler
val assembler = new VectorAssembler()
  .setInputCols(featureCols)
  .setOutputCol("features")

// Initialize the Linear Regression model
val lrModel = new LinearRegression()
  .setLabelCol("W_12_26") // Your target variable (Y)
  .setFeaturesCol("features") // The assembled feature vector (X)

// Build a pipeline to chain the assembler and model
val pipeline = new Pipeline().setStages(Array(assembler, lrModel))

// Split your data into training (70%) and test (30%) sets
val Array(trainingData, testData) = cleanedDF.randomSplit(Array(0.7, 0.3), seed = 1234)

// Train the model on the training data
val trainedModel = pipeline.fit(trainingData)

// Generate predictions on the test data
val predictions = trainedModel.transform(testData)

// View the actual labels vs. predicted values
predictions.select("W_12_26", "prediction").show()

// Evaluate the model's performance (e.g., calculate RMSE)
val evaluator = new RegressionEvaluator()
  .setLabelCol("W_12_26")
  .setPredictionCol("prediction")
  .setMetricName("rmse") // Options: rmse, mse, r2, mae

val rmse = evaluator.evaluate(predictions)
println(s"Root Mean Squared Error on Test Data: $rmse")

Full Python Code

from pyspark.ml.feature import VectorAssembler
from pyspark.ml.regression import LinearRegression
from pyspark.ml import Pipeline
from pyspark.ml.evaluation import RegressionEvaluator

# Define feature columns (exclude W_12_26)
feature_cols = [col for col in cleaned_df.columns if col != "W_12_26"]

# Create feature assembler
assembler = VectorAssembler(
    inputCols=feature_cols,
    outputCol="features"
)

# Initialize Linear Regression model
lr_model = LinearRegression(
    labelCol="W_12_26",
    featuresCol="features"
)

# Build pipeline
pipeline = Pipeline(stages=[assembler, lr_model])

# Split data into training and test sets
training_data, test_data = cleaned_df.randomSplit([0.7, 0.3], seed=1234)

# Train the model
trained_model = pipeline.fit(training_data)

# Generate predictions
predictions = trained_model.transform(test_data)

# View actual vs predicted values
predictions.select("W_12_26", "prediction").show()

# Evaluate model performance
evaluator = RegressionEvaluator(
    labelCol="W_12_26",
    predictionCol="prediction",
    metricName="rmse"
)

rmse = evaluator.evaluate(predictions)
print(f"Root Mean Squared Error on Test Data: {rmse}")

Key Notes for You as a Spark Newbie #

• DataSet vs DataFrame: Spark ML APIs are designed for DataFrames, so converting your DataSet to a DataFrame (with .toDF()) is the right move here.
• Feature Types: Make sure all your feature columns are numeric (double/float)—your data looks like it already is, but if you ever have categorical columns, you’ll need to encode them first (e.g., with OneHotEncoder).
• Tuning the Model: You can tweak model parameters to improve performance, like setMaxIter(100) (increase maximum iterations) or setRegParam(0.1) (add regularization to prevent overfitting).

内容的提问来源于stack exchange，提问作者Anna Klein